{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dependencies\n",
    "!pip install dspy-ai[chromadb] -Uqq\n",
    "!pip install termcolor -Uqq\n",
    "!pip install sqlalchemy -Uqq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AIM OF THE TUTORIAL\n",
    "\n",
    "* Build an end-to-end Text-to-SQL pipeline inspired from this [video](https://www.youtube.com/watch?v=L1o1VPVfbb0&pp=ygUYYWR2YW5jZWQgUkFHIGxsYW1hIGluZGV4) from Llama Index. In Llama Index, they used llama index Query Pipeline to build a Text-to-SQL pipeline. Here, we will build a Text-to-SQL pipeline based on our own dataset and from scratch. We will go from the dataset scraping, to building SQLlite database and using DSPy signatures to implementa a text-to-SQL pipeline\n",
    "\n",
    "## ABOUT THE DATASET\n",
    "* You can find the dataset [here](https://pages.stern.nyu.edu/~adamodar/New_Home_Page/datacurrent.html). The dataset has different industry based different financial metrics like WACC, tax rates, EBITDA, etc. There are multiple regions data `['US', 'Europe', 'Japan', 'AUS_NZ_CANADA', 'Emerging', 'China', 'India', 'Global']` where we have multiple tables for each region. There are nearly 250 tables with multiple columns in each table. We will build a text-to-SQL pipeline based on our own dataset and from scratch, starting from embedding tables schema and rows using ChromDB vector database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#imports\n",
    "from bs4 import BeautifulSoup \n",
    "import urllib.request\n",
    "import ssl\n",
    "from dotenv import load_dotenv\n",
    "import openai\n",
    "import os\n",
    "import requests\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import json\n",
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    ")\n",
    "import re\n",
    "from sqlalchemy import inspect\n",
    "import sqlalchemy\n",
    "from sqlalchemy import text \n",
    "import dspy\n",
    "from typing import List\n",
    "from termcolor import colored\n",
    "import chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dsp.modules.cache_utils import cache_turn_on\n",
    "\n",
    "cache_turn_on"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SCRAPING THE LINKS OF THE EXCEL FILES FROM THE WEBSITE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ssl._create_default_https_context = ssl._create_stdlib_context\n",
    "html_link = \"https://pages.stern.nyu.edu/~adamodar/New_Home_Page/datacurrent.html\"\n",
    "\n",
    "with urllib.request.urlopen(html_link) as url:\n",
    "    s = url.read()\n",
    "    # I'm guessing this would output the html source code ?\n",
    "    soup = BeautifulSoup(s,\"lxml\")\n",
    "\n",
    "html_table = soup.find_all(\"table\")\n",
    "req_table = html_table[1]\n",
    "hrefs_list = req_table.find_all('a')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "req_href = {\"US\":[],\"Europe\":[],\"Japan\":[],\"AUS_NZ_CANADA\":[],\"Emerging\":[],\"China\":[],\"India\":[],\"Global\":[]}\n",
    "\n",
    "for i in hrefs_list:\n",
    "    name = i.get_text().strip()\n",
    "    try:\n",
    "        href_attr = i['href']\n",
    "        # Only get the excel files\n",
    "        if href_attr.endswith('.xls'):\n",
    "            if \"US\" in name:\n",
    "                req_href[\"US\"].append(href_attr)\n",
    "            elif \"Europe\" in name:\n",
    "                req_href[\"Europe\"].append(href_attr)\n",
    "            elif \"Japan\" in name:\n",
    "                req_href[\"Japan\"].append(href_attr)\n",
    "            elif \"Aus\" in name:\n",
    "                req_href['AUS_NZ_CANADA'].append(href_attr)\n",
    "            elif \"Emerging\" in name:\n",
    "                req_href['Emerging'].append(href_attr)\n",
    "            elif \"China\" in name:\n",
    "                req_href['China'].append(href_attr)\n",
    "            elif \"India\" in name:\n",
    "                req_href['India'].append(href_attr)\n",
    "            elif \"Global\" in name: \n",
    "                req_href['Global'].append(href_attr)\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Download the excel files from the website and store it in a folder named DATA\n",
    "ssl._create_default_https_context = ssl._create_stdlib_context\n",
    "\n",
    "import os\n",
    "os.makedirs(\"DATA\",exist_ok=True)\n",
    "for country,excel_files in req_href.items():\n",
    "    country_path = os.path.join(\"DATA\",country) \n",
    "    os.makedirs(country_path,exist_ok=True)\n",
    "    for file in excel_files:\n",
    "        file_name = file.split(\"/\")[-1].split(\".\")[0]\n",
    "        full_file_name = os.path.join(country_path,f\"{file_name}.xls\")\n",
    "        resp = requests.get(file,verify=False)\n",
    "        output = open(full_file_name, 'wb')\n",
    "        output.write(resp.content)\n",
    "        output.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FOR Emerging WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR Europe WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR Global WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR AUS_NZ_CANADA WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR China WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR Japan WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n",
      "FOR US WE HAVE DIRECTORY LEN = 24 and ACTUAL LEN = 24\n",
      "FOR India WE HAVE DIRECTORY LEN = 29 and ACTUAL LEN = 29\n"
     ]
    }
   ],
   "source": [
    "# Sanity check\n",
    "for country in os.listdir(\"DATA\"):\n",
    "    dir_len = len(os.listdir(os.path.join(\"DATA\",country)))\n",
    "    country_len = len(req_href[country])\n",
    "    print(f'FOR {country} WE HAVE DIRECTORY LEN = {dir_len} and ACTUAL LEN = {country_len}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CLEANING THE DATASET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_excel = pd.ExcelFile(\"DATA2/US/capex.xls\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>End Game</th>\n",
       "      <th>To measure how much companies are reinvesting back into their long term assets, as a prelude to forecasting expected growth.</th>\n",
       "      <th>Unnamed: 2</th>\n",
       "      <th>Unnamed: 3</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Variable</td>\n",
       "      <td>How it is measured</td>\n",
       "      <td>What it measures</td>\n",
       "      <td>Units</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Capital Expenditures</td>\n",
       "      <td>Sum of the capital expenditures reported on st...</td>\n",
       "      <td>Gross investment in long term assets, at least...</td>\n",
       "      <td>$ millions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Depreciation</td>\n",
       "      <td>Sum of the depreciation reported on statement ...</td>\n",
       "      <td>Loss in value of assets, from use, as measured...</td>\n",
       "      <td>$ millions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Net Cap Ex</td>\n",
       "      <td>Sum of capital expenditures on statement of ca...</td>\n",
       "      <td>Net investment in long term assets, at least a...</td>\n",
       "      <td>$ millions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Net R&amp;D</td>\n",
       "      <td>Sum of R&amp;D reported as expense in most recent ...</td>\n",
       "      <td>Net investment in long term assets, expanded t...</td>\n",
       "      <td>$ millions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Acquisitions</td>\n",
       "      <td>Sum of acquisitions reported on statement of c...</td>\n",
       "      <td>Augments investment to include acquisition.</td>\n",
       "      <td>$ millions</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               End Game  \\\n",
       "0                   NaN   \n",
       "1                   NaN   \n",
       "2              Variable   \n",
       "3  Capital Expenditures   \n",
       "4          Depreciation   \n",
       "5            Net Cap Ex   \n",
       "6               Net R&D   \n",
       "7          Acquisitions   \n",
       "\n",
       "  To measure how much companies are reinvesting back into their long term assets, as a prelude to forecasting expected growth.  \\\n",
       "0                                                NaN                                                                             \n",
       "1                                                NaN                                                                             \n",
       "2                                 How it is measured                                                                             \n",
       "3  Sum of the capital expenditures reported on st...                                                                             \n",
       "4  Sum of the depreciation reported on statement ...                                                                             \n",
       "5  Sum of capital expenditures on statement of ca...                                                                             \n",
       "6  Sum of R&D reported as expense in most recent ...                                                                             \n",
       "7  Sum of acquisitions reported on statement of c...                                                                             \n",
       "\n",
       "                                          Unnamed: 2  Unnamed: 3  \n",
       "0                                                NaN         NaN  \n",
       "1                                                NaN         NaN  \n",
       "2                                   What it measures       Units  \n",
       "3  Gross investment in long term assets, at least...  $ millions  \n",
       "4  Loss in value of assets, from use, as measured...  $ millions  \n",
       "5  Net investment in long term assets, at least a...  $ millions  \n",
       "6  Net investment in long term assets, expanded t...  $ millions  \n",
       "7        Augments investment to include acquisition.  $ millions  "
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sn = 'Variables & FAQ'\n",
    "sample_excel.parse(sn).head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Date updated:</th>\n",
       "      <th>2024-01-05 00:00:00</th>\n",
       "      <th>Unnamed: 2</th>\n",
       "      <th>Unnamed: 3</th>\n",
       "      <th>Unnamed: 4</th>\n",
       "      <th>Unnamed: 5</th>\n",
       "      <th>Unnamed: 6</th>\n",
       "      <th>Unnamed: 7</th>\n",
       "      <th>Unnamed: 8</th>\n",
       "      <th>Unnamed: 9</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Created by:</td>\n",
       "      <td>Aswath Damodaran, adamodar@stern.nyu.edu</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What is this data?</td>\n",
       "      <td>Capital Expenditures, Acquisitions and R&amp;D and...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>US companies</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Home Page:</td>\n",
       "      <td>http://www.damodaran.com</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Data website:</td>\n",
       "      <td>https://pages.stern.nyu.edu/~adamodar/New_Home...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Companies in each industry:</td>\n",
       "      <td>https://pages.stern.nyu.edu/~adamodar/pc/datas...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Variable definitions:</td>\n",
       "      <td>https://pages.stern.nyu.edu/~adamodar/New_Home...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Industry Name</td>\n",
       "      <td>Number of Firms</td>\n",
       "      <td>Capital Expenditures (US $ millions)</td>\n",
       "      <td>Depreciation &amp; Amort ((US $ millions)</td>\n",
       "      <td>Cap Ex/Deprecn</td>\n",
       "      <td>Acquisitions (US $ millions)</td>\n",
       "      <td>Net R&amp;D (US $ millions)</td>\n",
       "      <td>Net Cap Ex/Sales</td>\n",
       "      <td>Net Cap Ex/ EBIT (1-t)</td>\n",
       "      <td>Sales/ Invested Capital (LTM)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Advertising</td>\n",
       "      <td>57</td>\n",
       "      <td>775.729</td>\n",
       "      <td>1887.338</td>\n",
       "      <td>0.411018</td>\n",
       "      <td>322.559</td>\n",
       "      <td>75.08</td>\n",
       "      <td>-0.016886</td>\n",
       "      <td>-0.21812</td>\n",
       "      <td>3.283403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Aerospace/Defense</td>\n",
       "      <td>70</td>\n",
       "      <td>10982.128</td>\n",
       "      <td>13311.598</td>\n",
       "      <td>0.825004</td>\n",
       "      <td>10344.75</td>\n",
       "      <td>830.4634</td>\n",
       "      <td>0.022829</td>\n",
       "      <td>0.318077</td>\n",
       "      <td>1.98434</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Air Transport</td>\n",
       "      <td>25</td>\n",
       "      <td>25559.725</td>\n",
       "      <td>10609.948</td>\n",
       "      <td>2.409034</td>\n",
       "      <td>368.21</td>\n",
       "      <td>73.5486</td>\n",
       "      <td>0.067203</td>\n",
       "      <td>1.355173</td>\n",
       "      <td>1.77732</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 Date updated:  \\\n",
       "0                  Created by:   \n",
       "1           What is this data?   \n",
       "2                   Home Page:   \n",
       "3                Data website:   \n",
       "4  Companies in each industry:   \n",
       "5        Variable definitions:   \n",
       "6                Industry Name   \n",
       "7                  Advertising   \n",
       "8            Aerospace/Defense   \n",
       "9                Air Transport   \n",
       "\n",
       "                                 2024-01-05 00:00:00  \\\n",
       "0           Aswath Damodaran, adamodar@stern.nyu.edu   \n",
       "1  Capital Expenditures, Acquisitions and R&D and...   \n",
       "2                           http://www.damodaran.com   \n",
       "3  https://pages.stern.nyu.edu/~adamodar/New_Home...   \n",
       "4  https://pages.stern.nyu.edu/~adamodar/pc/datas...   \n",
       "5  https://pages.stern.nyu.edu/~adamodar/New_Home...   \n",
       "6                                    Number of Firms   \n",
       "7                                                 57   \n",
       "8                                                 70   \n",
       "9                                                 25   \n",
       "\n",
       "                             Unnamed: 2  \\\n",
       "0                                   NaN   \n",
       "1                                   NaN   \n",
       "2                                   NaN   \n",
       "3                                   NaN   \n",
       "4                                   NaN   \n",
       "5                                   NaN   \n",
       "6  Capital Expenditures (US $ millions)   \n",
       "7                               775.729   \n",
       "8                             10982.128   \n",
       "9                             25559.725   \n",
       "\n",
       "                              Unnamed: 3      Unnamed: 4  \\\n",
       "0                                    NaN             NaN   \n",
       "1                                    NaN             NaN   \n",
       "2                                    NaN             NaN   \n",
       "3                                    NaN             NaN   \n",
       "4                                    NaN             NaN   \n",
       "5                                    NaN             NaN   \n",
       "6  Depreciation & Amort ((US $ millions)  Cap Ex/Deprecn   \n",
       "7                               1887.338        0.411018   \n",
       "8                              13311.598        0.825004   \n",
       "9                              10609.948        2.409034   \n",
       "\n",
       "                     Unnamed: 5               Unnamed: 6        Unnamed: 7  \\\n",
       "0                           NaN                      NaN               NaN   \n",
       "1                  US companies                      NaN               NaN   \n",
       "2                           NaN                      NaN               NaN   \n",
       "3                           NaN                      NaN               NaN   \n",
       "4                           NaN                      NaN               NaN   \n",
       "5                           NaN                      NaN               NaN   \n",
       "6  Acquisitions (US $ millions)  Net R&D (US $ millions)  Net Cap Ex/Sales   \n",
       "7                       322.559                    75.08         -0.016886   \n",
       "8                      10344.75                 830.4634          0.022829   \n",
       "9                        368.21                  73.5486          0.067203   \n",
       "\n",
       "               Unnamed: 8                     Unnamed: 9  \n",
       "0                     NaN                            NaN  \n",
       "1                     NaN                            NaN  \n",
       "2                     NaN                            NaN  \n",
       "3                     NaN                            NaN  \n",
       "4                     NaN                            NaN  \n",
       "5                     NaN                            NaN  \n",
       "6  Net Cap Ex/ EBIT (1-t)  Sales/ Invested Capital (LTM)  \n",
       "7                -0.21812                       3.283403  \n",
       "8                0.318077                        1.98434  \n",
       "9                1.355173                        1.77732  "
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_excel.parse(sample_excel.sheet_names[1]).head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* In the dataset above, there are two sheets. The first sheet is the variables and summary, and the second sheet is the table with the data. \n",
    "* We will clean the first sheet to get the table name and summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Emerging\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 86.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Europe\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 94.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 102.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUS_NZ_CANADA\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 104.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "China\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 103.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Japan\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 66.34it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "US\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 24/24 [00:00<00:00, 102.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "India\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29/29 [00:00<00:00, 100.35it/s]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "pd.set_option('display.max_rows', 50)\n",
    "\n",
    "def sanitize_column_name(col_name):\n",
    "    # Remove special characters and replace spaces with underscores\n",
    "    return re.sub(r\"\\W+\", \"_\", col_name)\n",
    "\n",
    "dir = \"DATA\"\n",
    "processed_dir = \"Processed Data\"\n",
    "all_infos_dict = []\n",
    "os.makedirs(processed_dir,exist_ok=True)\n",
    "for country in os.listdir(dir):\n",
    "    print(country)\n",
    "    file_name = os.path.join(dir,country)\n",
    "    os.makedirs(os.path.join(processed_dir,country),exist_ok=True)\n",
    "    os.makedirs(file_name,exist_ok=True)\n",
    "    for excel_file in tqdm(os.listdir(file_name)):\n",
    "        full_file_name = os.path.join(file_name,excel_file)\n",
    "        xls = pd.ExcelFile(full_file_name)\n",
    "        sns = xls.sheet_names\n",
    "        for sheet_name in sns:\n",
    "            if \"Var\" in sheet_name or \"var\" in sheet_name:\n",
    "                info_df = xls.parse(sheet_name)\n",
    "                info_df.dropna(how=\"all\",inplace=True)\n",
    "                info_dict = {}\n",
    "                for cols in info_df.columns:\n",
    "                    if \"End\" not in cols and 'Unnamed' not in cols:\n",
    "                        info_dict['Summary'] = cols\n",
    "                info_dict['Vars'] = info_df.values[1:].tolist()\n",
    "                all_infos_dict.append(info_dict)\n",
    "            elif \"Industry\" in sheet_name or \"industry\" in sheet_name:\n",
    "                data_df = xls.parse(sheet_name)\n",
    "        try:\n",
    "            data_df.dropna(axis=1,thresh=5,inplace=True)\n",
    "            data_df.dropna(inplace=True)\n",
    "            new_header = data_df.iloc[0] #grab the first row for the header\n",
    "        except:\n",
    "            print(full_file_name)\n",
    "            print(data_df)\n",
    "        data_df = data_df[1:] #take the data less the header row\n",
    "        data_df.reset_index(inplace=True,drop=True)\n",
    "        new_header = [sanitize_column_name(str(col)) for col in new_header]\n",
    "        data_df.columns = new_header #set the header row as the df header\n",
    "        save_name = full_file_name.split(\".\")[0].split(\"/\")[-1]\n",
    "        save_file_path = os.path.join(os.path.join(processed_dir,country),save_name)\n",
    "        data_df.to_csv(save_file_path+\".csv\",index=False)\n",
    "        with open(save_file_path+\".json\", \"w\") as outfile: \n",
    "            json.dump(info_dict, outfile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A SAMPLE METADATA JSON"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Summary': 'Measures of accounting returns, to all claim holders (and from operations)',\n",
       " 'Vars': [['Number of firms',\n",
       "   'Number of firms in the industry grouping.',\n",
       "   'Law of large numbers?'],\n",
       "  ['R&D Capitalized',\n",
       "   'My estimate of R&D capitalization, based upon a 5-year straight line amortization period, aggregated across firms in the group',\n",
       "   'Capitalized value of R&D gets added on to book equity and to invested capital'],\n",
       "  ['Capitalized R&D as percent of invested capital',\n",
       "   'My R&D capitalization estimate, as a percent of invested capital including that number, based upon aggregated values across firms',\n",
       "   'Magnitude of investment in R&D, relative to investment in more traditional capital expenditures.'],\n",
       "  ['R&D -LTM',\n",
       "   'Aggregated R&D expenses across the last twelve months, across companies in the group.',\n",
       "   'Spending on R&D in most recent year'],\n",
       "  ['R&D: Years minus 1 to minus 5',\n",
       "   'Aggregated R&D expenses for each of the previous five years, across companies in the group.',\n",
       "   'Sepnding on R&D over last five years']]}"
      ]
     },
     "execution_count": 196,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_infos_dict[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DATAFRAME AFTER PREPROCESSING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Industry_Name</th>\n",
       "      <th>Number_of_Firms</th>\n",
       "      <th>Capital_Expenditures_(US_$_millions)</th>\n",
       "      <th>Depreciation_&amp;_Amort_((US_$_millions)</th>\n",
       "      <th>Cap_Ex/Deprecn</th>\n",
       "      <th>Acquisitions_(US_$_millions)</th>\n",
       "      <th>Net_R&amp;D_(US_$_millions)</th>\n",
       "      <th>Net_Cap_Ex/Sales</th>\n",
       "      <th>Net_Cap_Ex/_EBIT_(1-t)</th>\n",
       "      <th>Sales/_Invested_Capital_(LTM)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Advertising</td>\n",
       "      <td>57</td>\n",
       "      <td>775.729</td>\n",
       "      <td>1887.338</td>\n",
       "      <td>0.411018</td>\n",
       "      <td>322.559</td>\n",
       "      <td>75.0800</td>\n",
       "      <td>-0.016886</td>\n",
       "      <td>-0.218120</td>\n",
       "      <td>3.283403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Aerospace/Defense</td>\n",
       "      <td>70</td>\n",
       "      <td>10982.128</td>\n",
       "      <td>13311.598</td>\n",
       "      <td>0.825004</td>\n",
       "      <td>10344.750</td>\n",
       "      <td>830.4634</td>\n",
       "      <td>0.022829</td>\n",
       "      <td>0.318077</td>\n",
       "      <td>1.984340</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Air Transport</td>\n",
       "      <td>25</td>\n",
       "      <td>25559.725</td>\n",
       "      <td>10609.948</td>\n",
       "      <td>2.409034</td>\n",
       "      <td>368.210</td>\n",
       "      <td>73.5486</td>\n",
       "      <td>0.067203</td>\n",
       "      <td>1.355173</td>\n",
       "      <td>1.777320</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Apparel</td>\n",
       "      <td>38</td>\n",
       "      <td>1730.981</td>\n",
       "      <td>1386.480</td>\n",
       "      <td>1.248472</td>\n",
       "      <td>38.739</td>\n",
       "      <td>0.9474</td>\n",
       "      <td>0.005365</td>\n",
       "      <td>0.072557</td>\n",
       "      <td>1.773076</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Auto &amp; Truck</td>\n",
       "      <td>34</td>\n",
       "      <td>29899.486</td>\n",
       "      <td>18677.668</td>\n",
       "      <td>1.600815</td>\n",
       "      <td>193.220</td>\n",
       "      <td>983.0696</td>\n",
       "      <td>0.026577</td>\n",
       "      <td>0.675918</td>\n",
       "      <td>1.048732</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       Industry_Name  Number_of_Firms  Capital_Expenditures_(US_$_millions)  \\\n",
       "0        Advertising               57                               775.729   \n",
       "1  Aerospace/Defense               70                             10982.128   \n",
       "2      Air Transport               25                             25559.725   \n",
       "3            Apparel               38                              1730.981   \n",
       "4       Auto & Truck               34                             29899.486   \n",
       "\n",
       "   Depreciation_&_Amort_((US_$_millions)  Cap_Ex/Deprecn  \\\n",
       "0                               1887.338        0.411018   \n",
       "1                              13311.598        0.825004   \n",
       "2                              10609.948        2.409034   \n",
       "3                               1386.480        1.248472   \n",
       "4                              18677.668        1.600815   \n",
       "\n",
       "   Acquisitions_(US_$_millions)  Net_R&D_(US_$_millions)  Net_Cap_Ex/Sales  \\\n",
       "0                       322.559                  75.0800         -0.016886   \n",
       "1                     10344.750                 830.4634          0.022829   \n",
       "2                       368.210                  73.5486          0.067203   \n",
       "3                        38.739                   0.9474          0.005365   \n",
       "4                       193.220                 983.0696          0.026577   \n",
       "\n",
       "   Net_Cap_Ex/_EBIT_(1-t)  Sales/_Invested_Capital_(LTM)  \n",
       "0               -0.218120                       3.283403  \n",
       "1                0.318077                       1.984340  \n",
       "2                1.355173                       1.777320  \n",
       "3                0.072557                       1.773076  \n",
       "4                0.675918                       1.048732  "
      ]
     },
     "execution_count": 197,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"Processed Data/US/capex.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BUILD TABLE NAMES AND METADATA\n",
    "\n",
    "* Here we use a DSPy signature given the first 10 rows of the dataframe, we generate the table name and table explanation. It will help us to dynamically select the correct table based on the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "',Industry_Name,Number_of_Firms,Capital_Expenditures_(US_$_millions),Depreciation_&_Amort_((US_$_millions),Cap_Ex/Deprecn,Acquisitions_(US_$_millions),Net_R&D_(US_$_millions),Net_Cap_Ex/Sales,Net_Cap_Ex/_EBIT_(1-t),Sales/_Invested_Capital_(LTM)\\n0,Advertising,57,775.729,1887.338,0.4110175283918408,322.559,75.07999999999993,-0.0168860862168837,-0.2181203933865563,3.283403041333076\\n1,Aerospace/Defense,70,10982.128,13311.597999999998,0.8250044810547916,10344.749999999998,830.4633999999933,0.0228292321038578,0.318077061754813,1.9843395804666923\\n2,Air Transport,25,25559.725,10609.948,2.4090339556800844,368.21,73.5486000000019,0.0672027546276203,1.3551729988499104,1.7773199625336142\\n3,Apparel,38,1730.9809999999998,1386.4799999999996,1.2484716692631703,38.739,0.9473999999991064,0.0053650838666078,0.0725567415076527,1.7730762352050575\\n4,Auto & Truck,34,29899.486,18677.668,1.6008147269776931,193.22,983.0695999999988,0.0265766732899557,0.6759175844836917,1.048731678622114\\n5,Auto Parts,39,3565.4210000000007,2034.471,1.7525051966825778,1061.7600000000002,454.9352000000008,0.0316936036463534,0.8398185079860373,2.004790963758861\\n6,Beverage (Alcoholic),19,2244.8160000000007,864.0629999999999,2.59797723082692,1370.0,0.0,0.0939030584900288,0.6367443614850201,0.8968231153352131\\n7,Beverage (Soft),29,7926.687000000001,4543.389,1.7446639501922463,674.6669999999999,7.705600000001141,0.0238699281532688,0.1517545513863501,1.6931990505110932\\n8,Broadcasting,22,1837.548,3387.6400000000003,0.5424271764414165,43.856,0.2399999999997817,-0.0206121524427381,-0.2104150803584717,1.1159893806472496\\n9,Brokerage & Investment Banking,27,7783.945999999999,4646.382,1.675270350134793,876.0999999999999,-189.38079999999945,0.0167661627913481,3.4298253908840737,0.2794992243703733\\n'"
      ]
     },
     "execution_count": 199,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head(10).to_csv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "openai.api_key = os.environ['OPENAI_API_KEY']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "turbo = dspy.OpenAI(model='gpt-3.5-turbo-instruct', max_tokens=250)\n",
    "dspy.settings.configure(lm=turbo)\n",
    "\n",
    "class SQLTableMetadata(dspy.Signature):\n",
    "    \"\"\"Give a suitable table name and description about the given table\"\"\"\n",
    "    pandas_dataframe_str = dspy.InputField(desc=\"First 10 rows of a pandas dataframe delimited by newline character\")\n",
    "    table_name = dspy.OutputField(desc=\"suitable table name\")\n",
    "    table_summary = dspy.OutputField(desc=\"a summary about the table\")\n",
    "\n",
    "class CoT(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.prog = dspy.ChainOfThought(SQLTableMetadata)\n",
    "    \n",
    "    def forward(self, pandas_dataframe_str):\n",
    "        return self.prog(pandas_dataframe_str=pandas_dataframe_str)\n",
    "\n",
    "cot = CoT()\n",
    "\n",
    "# cot(pandas_dataframe_str = df.head(10).to_csv())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_dir = \"Processed Data\"\n",
    "dfs_str = []\n",
    "for country in os.listdir(processed_dir):\n",
    "    country_folder = os.path.join(processed_dir,country)\n",
    "    # print(f\"{country}\")\n",
    "    for files in tqdm(os.listdir(country_folder),desc=f\"Building the summary and name for {country}\"):\n",
    "        if files.endswith(\".csv\"):\n",
    "            file_name = files.split(\".\")[0]\n",
    "            csv_file_path = os.path.join(country_folder,files)\n",
    "            df = pd.read_csv(csv_file_path,index_col=False)\n",
    "            json_file_path = os.path.join(country_folder,f\"{file_name}.json\")\n",
    "            with open(json_file_path,'r') as f:\n",
    "                data = json.loads(f.read())\n",
    "            if 'table_name' in data and 'table_summary' in data:\n",
    "                # if data['table_name'] == \"\" or data['table_summary'] == \"\":\n",
    "                if data['table_summary'] == \"\":\n",
    "                    pass\n",
    "                else:\n",
    "                    continue\n",
    "            dfs_str.append(df.head(10).to_csv())\n",
    "            table_preds = cot(pandas_dataframe_str = df.head(10).to_csv())\n",
    "            data['table_name'] = table_preds.table_name\n",
    "            data['table_summary'] = table_preds.table_summary\n",
    "            with open(json_file_path,'w') as f:\n",
    "                json.dump(data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NEXT TASKS\n",
    "1. Build database with each region for each table\n",
    "2. Embed the table summary and table SCHEMA. Also, embed the table rows\n",
    "3. Retrieval at table level and embed the rows to retrieve relevant rows from the retrieved schema of table\n",
    "4. Text-to-SQL pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BUILD THE SQLITE DATABASE FROM THE CSV FILES\n",
    "\n",
    "It was taken from the [tutorial](https://docs.llamaindex.ai/en/stable/examples/pipeline/query_pipeline_sql/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create a sanitized column name\n",
    "def sanitize_column_name(col_name):\n",
    "    # Remove special characters and replace spaces with underscores\n",
    "    return re.sub(r\"\\W+\", \"_\", col_name)\n",
    "\n",
    "\n",
    "# Function to create a table from a DataFrame using SQLAlchemy\n",
    "def create_table_from_dataframe(\n",
    "    df: pd.DataFrame, table_name: str, engine, metadata_obj\n",
    "):\n",
    "    # Sanitize column names\n",
    "    sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}\n",
    "    df = df.rename(columns=sanitized_columns)\n",
    "\n",
    "    # Dynamically create columns based on DataFrame columns and data types\n",
    "    columns = [\n",
    "        Column(col, String if dtype == \"object\" else Integer)\n",
    "        for col, dtype in zip(df.columns, df.dtypes)\n",
    "    ]\n",
    "\n",
    "    # Create a table with the defined columns\n",
    "    table = Table(table_name, metadata_obj, *columns)\n",
    "\n",
    "    # Create the table in the database\n",
    "    metadata_obj.create_all(engine)\n",
    "\n",
    "    # Insert data from DataFrame into the table\n",
    "    with engine.connect() as conn:\n",
    "        for _, row in df.iterrows():\n",
    "            insert_stmt = table.insert().values(**row.to_dict())\n",
    "            conn.execute(insert_stmt)\n",
    "        conn.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DATABASE CREATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_dir  = \"Processed Data\"\n",
    "def sqlalchemy_engine(region:str):\n",
    "    \"\"\"Create a SQLAlchemy engine for the given region\"\"\"\n",
    "    assert region in os.listdir(processed_dir), f\"{region} is not a valid region from {os.listdir(processed_dir)}\"\n",
    "    # Create a SQLAlchemy database for each region\n",
    "    engine = create_engine(f\"sqlite:///{region}.db\")\n",
    "    metadata_obj = MetaData()\n",
    "    region_path = os.path.join(processed_dir,region)\n",
    "    dfs = []\n",
    "    for dataframes_path in os.listdir(region_path):\n",
    "        if dataframes_path.endswith(\".csv\"):\n",
    "            df = pd.read_csv(os.path.join(region_path,dataframes_path),index_col=False)\n",
    "            dfs.append((dataframes_path,df))\n",
    "    pbar = tqdm(total=len(dfs),desc=f\"Creating tables for {region}\")\n",
    "    for _, df_table_name in enumerate(dfs):\n",
    "        table_name = df_table_name[0]\n",
    "        table_name = table_name.split(\".\")[0]\n",
    "        df = df_table_name[1]\n",
    "        # print(f\"Creating table: {table_name}\")\n",
    "        create_table_from_dataframe(df,table_name, engine, metadata_obj)\n",
    "        # print(f\"Done creating table for: {table_name}\")\n",
    "        pbar.update(1)\n",
    "    return engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating tables for US:   0%|          | 0/24 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating tables for US: 100%|██████████| 24/24 [00:00<00:00, 32.16it/s]\n",
      "Creating tables for India: 100%|██████████| 29/29 [00:00<00:00, 33.36it/s]\n",
      "Creating tables for China: 100%|██████████| 29/29 [00:01<00:00, 28.46it/s]\n",
      "Creating tables for Europe: 100%|██████████| 29/29 [00:00<00:00, 30.79it/s]\n",
      "Creating tables for Global: 100%|██████████| 29/29 [00:00<00:00, 30.62it/s]\n",
      "Creating tables for AUS_NZ_CANADA: 100%|██████████| 29/29 [00:00<00:00, 29.32it/s]\n",
      "Creating tables for Japan: 100%|██████████| 29/29 [00:00<00:00, 33.82it/s]\n",
      "Creating tables for Emerging: 100%|██████████| 29/29 [00:00<00:00, 31.91it/s]\n"
     ]
    }
   ],
   "source": [
    "us_engine = sqlalchemy_engine(\"US\")\n",
    "india_engine = sqlalchemy_engine(\"India\")\n",
    "china_engine = sqlalchemy_engine(\"China\")\n",
    "europe_engine = sqlalchemy_engine(\"Europe\")\n",
    "global_engine = sqlalchemy_engine(\"Global\")\n",
    "aus_nz_canada_engine = sqlalchemy_engine(\"AUS_NZ_CANADA\")\n",
    "japan_engine = sqlalchemy_engine(\"Japan\")\n",
    "emerging_engine = sqlalchemy_engine(\"Emerging\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_infos(sql_engine:sqlalchemy.engine.base.Engine,region:str):\n",
    "    \"\"\"Get all the tables info in the database based on the given region\"\"\"\n",
    "    inspector = inspect(sql_engine)\n",
    "    table_names = inspector.get_table_names()\n",
    "    table_infos_dict = {tb: [] for tb in table_names}\n",
    "    for tb in table_names:\n",
    "        column_dict = inspector.get_columns(tb)\n",
    "        schema_str = \"\"\n",
    "        primary_keys = []\n",
    "        for col in column_dict:\n",
    "            schema_str += f\"{col['name']} ({col['type']}), \"\n",
    "            if col[\"primary_key\"] not in primary_keys:\n",
    "                primary_keys.append(col[\"name\"])\n",
    "        with open(os.path.join(processed_dir,region,f\"{tb}.json\")) as f:\n",
    "            table_info = json.loads(f.read())\n",
    "        table_infos_dict[tb] = [\n",
    "            {\n",
    "                \"table_info\": f\"Table {tb} has columns: {schema_str[:-2]}\",\n",
    "                \"table_summary\": f'{table_info[\"Summary\"]}. {table_info[\"table_summary\"]}. ',\n",
    "            }\n",
    "        ]\n",
    "    return table_infos_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "us_tb_dict = get_table_infos(us_engine,\"US\")\n",
    "india_tb_dict = get_table_infos(india_engine,\"India\")\n",
    "china_tb_dict = get_table_infos(china_engine,\"China\")\n",
    "europe_tb_dict = get_table_infos(europe_engine,\"Europe\")\n",
    "global_tb_dict = get_table_infos(global_engine,\"Global\")\n",
    "aus_nz_canada_tb_dict = get_table_infos(aus_nz_canada_engine,\"AUS_NZ_CANADA\")\n",
    "japan_tb_dict = get_table_infos(japan_engine,\"Japan\")\n",
    "emerging_tb_dict = get_table_infos(emerging_engine,\"Emerging\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'table_info': 'Table DollarUS has columns: Industry_Name (VARCHAR), Number_of_firms (INTEGER), Average_Company_Age_years_ (INTEGER), Market_Cap_millions_ (INTEGER), Book_Equity_millions_ (INTEGER), Enteprise_Value_millions_ (INTEGER), Invested_Capital_millions_ (INTEGER), Total_Debt_including_leases_millions_ (INTEGER), Revenues_millions_ (INTEGER), Gross_Profit_millions_ (INTEGER), EBITDA_millions_ (INTEGER), EBIT_Operating_Income_millions_ (INTEGER), Net_Income_millions_ (INTEGER)',\n",
       "  'table_summary': 'To report aggregated dollar value of key operating and marker numbers, by industry group, in millions of US $.. This table contains financial information about various industries such as their names, number of firms, average company age, market cap, book equity, enterprise value, invested capital, total debt, revenues, gross profit, EBITDA, EBIT, and net income. The data is delimited by a newline character and can be used for comparison and analysis of the financial performance of different industries.. '}]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "us_tb_dict['DollarUS']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EMBEDDINGS\n",
    "\n",
    "1. Embed the table summary and table SCHEMA to get the table that the user is looking for\n",
    "2. Embed the table rows for each table, so as to get relevant rows from the retrieved table "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EMBED THE TABLE SUMMARY AND TABLE SCHEMA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "import chromadb.utils.embedding_functions as embedding_functions\n",
    "from chromadb.utils.batch_utils import create_batches\n",
    "\n",
    "load_dotenv(override=True)\n",
    "emb_fn = embedding_functions.OpenAIEmbeddingFunction(\n",
    "                api_key=os.environ['OPENAI_API_KEY'],\n",
    "                model_name=\"text-embedding-3-small\")\n",
    "# EMBEDDING_MODEL = \"mixedbread-ai/mxbai-embed-large-v1\"\n",
    "# emb_fn = embedding_functions.HuggingFaceEmbeddingFunction(model_name=EMBEDDING_MODEL,api_key=os.environ[\"HF_API_KEY\"])\n",
    "def embed_table_info(region:str,tb_dict):\n",
    "    \"\"\"Embed the table summary and table SCHEMA to get the table that the user is looking for\"\"\"\n",
    "    client = chromadb.PersistentClient(path=f\"{region}_TABLE\")\n",
    "\n",
    "    table_collection = client.create_collection(name=\"table\",embedding_function=emb_fn)\n",
    "\n",
    "    table_docs = []\n",
    "    table_metadata = []\n",
    "\n",
    "\n",
    "    for table_name,table_data in tb_dict.items():\n",
    "        table_docs.append(table_data[0]['table_info'] + \". \" + table_data[0]['table_summary'])\n",
    "        table_metadata.append({\"table_name\":table_name,'table_metadata':table_data[0]['table_info']})\n",
    "    table_ids = [f\"id{i}\" for i in range(len(table_docs))]\n",
    "    assert len(table_docs) == len(table_metadata)\n",
    "    print(len(table_docs),len(table_metadata))\n",
    "    # Create a batch of data to be sent to OpenAI Embedding API\n",
    "    batches = create_batches(api=client,ids=table_ids, documents=table_docs, metadatas=table_metadata)\n",
    "    for batch in tqdm(batches,desc=\"Embedding table info\"):\n",
    "        table_collection.add(ids=batch[0],\n",
    "                    documents=batch[3],\n",
    "                    metadatas=batch[2])\n",
    "\n",
    "# embed_table_info(\"US\",us_tb_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## For some strange reason, the `create_batches` was not batching the below documents, hence I had to do it manually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def embed_rows(region:str,batch_size:int=24):\n",
    "    client = chromadb.PersistentClient(path=f\"{region}_TABLE\")\n",
    "    # client.delete_collection(name=\"rows\")\n",
    "    rows_collection = client.create_collection(name=\"rows\",embedding_function=emb_fn)\n",
    "\n",
    "    rows_docs = []\n",
    "    rows_metadata = []\n",
    "    region_path = os.path.join(processed_dir,region)\n",
    "    for df_path in os.listdir(region_path):\n",
    "        df_full_path = os.path.join(region_path,df_path)\n",
    "        df = pd.read_csv(df_full_path,index_col=False)\n",
    "        for idx,row in df.iterrows():\n",
    "            row_str = \"\"\n",
    "            full_rows = []\n",
    "            for rv in row.values:\n",
    "                if isinstance(rv,str):\n",
    "                    row_str+= rv + \", \"\n",
    "                full_rows.append(str(rv))\n",
    "                row_str = row_str.replace('\"',\"\")\n",
    "                # row_str = row_str.replace(\"'\",'\"')\n",
    "            full_rows_str = \", \".join(full_rows)[:-2]\n",
    "            full_rows_str = full_rows_str.replace('\"',\"\")\n",
    "            rows_docs.append(row_str[:-2])\n",
    "            rows_metadata.append({\"table_name\":df_path.split(\".\")[0],\"region\":region,\"index\":idx,\"full_rows\":full_rows_str})\n",
    "    row_ids = [f\"id{i}\" for i in range(len(rows_docs))]\n",
    "    # print(len(rows_docs),len(rows_metadata))\n",
    "    assert len(rows_docs) == len(rows_metadata) == len(row_ids)\n",
    "    # return rows_docs,rows_metadata,row_ids\n",
    "    for start in tqdm(range(0,len(rows_docs),batch_size),desc=\"Embedding rows\"):\n",
    "        end = min(start+batch_size,len(rows_docs))\n",
    "        batch_ids = row_ids[start:end]\n",
    "        batch_rows = rows_docs[start:end]\n",
    "        batch_metadatas = rows_metadata[start:end]\n",
    "        rows_collection.add(ids=batch_ids,\n",
    "                    documents=batch_rows,\n",
    "                    metadatas=batch_metadatas)\n",
    "    # return batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:04<00:00,  2.44s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"US\"\n",
    "embed_table_info(region,us_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:05<00:00,  2.61s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"India\"\n",
    "embed_table_info(region,india_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:25<00:00, 12.86s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"China\"\n",
    "embed_table_info(region,china_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 3/3 [00:26<00:00,  8.89s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"Europe\"\n",
    "embed_table_info(region,europe_tb_dict)\n",
    "embed_rows(region,1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:25<00:00, 12.69s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"Global\"\n",
    "embed_table_info(region,global_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:25<00:00, 12.61s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"Emerging\"\n",
    "embed_table_info(region,emerging_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:45<00:00, 22.99s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"Japan\"\n",
    "embed_table_info(region,japan_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding rows: 100%|██████████| 2/2 [00:25<00:00, 12.68s/it]\n"
     ]
    }
   ],
   "source": [
    "region = \"AUS_NZ_CANADA\"\n",
    "embed_table_info(region,aus_nz_canada_tb_dict)\n",
    "embed_rows(region,2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TEXT-TO-SQL PIPELINE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LOAD DATABASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_dict = {\n",
    "    \"US\":us_engine,\n",
    "    \"India\":india_engine,\n",
    "    \"China\":china_engine,\n",
    "    \"Europe\":europe_engine,\n",
    "    \"Global\":global_engine,\n",
    "    \"AUS_NZ_CANADA\":aus_nz_canada_engine,\n",
    "    \"Japan\":japan_engine,\n",
    "    \"Emerging\":emerging_engine,\n",
    "}\n",
    "\n",
    "def get_collections_db(region:str):\n",
    "    # Get the database for the given region, table collection and row collection\n",
    "    client = chromadb.PersistentClient(path=f\"{region}_TABLE\")\n",
    "    table_collection = client.get_collection(name=\"table\",embedding_function=emb_fn)\n",
    "    row_collection = client.get_collection(name=\"rows\",embedding_function=emb_fn)\n",
    "    return [db_dict[region],table_collection,row_collection]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_collection_dict = {\n",
    "    \"US\":get_collections_db(\"US\"),\n",
    "    \"India\":get_collections_db(\"India\"),\n",
    "    \"China\":get_collections_db(\"China\"),\n",
    "    \"Europe\":get_collections_db(\"Europe\"),\n",
    "    \"Global\":get_collections_db(\"Global\"),\n",
    "    \"AUS_NZ_CANADA\":get_collections_db(\"AUS_NZ_CANADA\"),\n",
    "    \"Japan\":get_collections_db(\"Japan\"),\n",
    "    \"Emerging\":get_collections_db(\"Emerging\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "text_to_sql = dspy.OpenAI(model='gpt-3.5-turbo', max_tokens=1024)\n",
    "sql_to_answer = dspy.OpenAI(model='gpt-3.5-turbo',max_tokens=1024)\n",
    "\n",
    "# DSPy signature for converting text to SQL query\n",
    "class TextToSQLAnswer(dspy.Signature):\n",
    "    \"\"\"Convert natural language text to SQL using suitable schema(s) from multiple schema choices\"\"\"\n",
    "\n",
    "    question:str = dspy.InputField(desc=\"natural language input which will be converted to SQL\")\n",
    "    relevant_table_schemas_rows:str = dspy.InputField(desc=\"Multiple possible tables which has table name and corresponding columns, along with relevant rows from the table (values in the same order as columns above)\")\n",
    "    sql:str = dspy.OutputField(desc=\"Generate syntactically correct sqlite query with correct column names using suitable tables(s) and its rows.\\n Don't forget to add distinct.\\n Please rename the returned columns into suitable names.\\n DON'T OUTPUT anything else other than the sqlite query\")\n",
    "\n",
    "# DSPy signature for converting SQL query and question to natural language text\n",
    "class SQLReturnToAnswer(dspy.Signature):\n",
    "    \"\"\"Answer the question using the rows from the SQL query\"\"\"\n",
    "    question:str = dspy.InputField()\n",
    "    sql:str = dspy.InputField(desc=\"sqlite query that generated the rows\")\n",
    "    relevant_rows:str = dspy.InputField(desc=\"relevant rows to answer the question\")\n",
    "    answer:str = dspy.OutputField(desc=\"answer to the question using relevant rows and the sql query\")\n",
    "\n",
    "# If there is an SQLError, then rectify the error by trying again\n",
    "class SQLRectifier(dspy.Signature):\n",
    "    \"\"\"Correct the SQL query to resolve the error using the proper table names, columns and rows\"\"\"  \n",
    "    input_sql:str = dspy.InputField(desc=\"sqlite query that needs to be fixed\")\n",
    "    error_str: str = dspy.InputField(desc=\"error that needs to be resolved\")\n",
    "    relevant_table_schemas_rows:str = dspy.InputField(desc=\"Multiple possible tables which has table name and corresponding columns, along with relevant rows from the table (values in the same order as columns above)\")\n",
    "    sql:str = dspy.OutputField(desc=\"corrected sqlite query to resolve the error and remove and any invalid syntax in the query.\\n Don't output anything else other than the sqlite query\")\n",
    "\n",
    "dspy.settings.configure(lm=text_to_sql)\n",
    "\n",
    "# Filter out the SQL Query\n",
    "def process_sql_str(sql_str:str):\n",
    "    sql_str = sql_str.replace(\"```\",\"\")\n",
    "    sql_str = sql_str.replace(\"sql\",\"\")\n",
    "    sql_str = sql_str.strip()\n",
    "    return sql_str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p align=\"center\">\n",
    "  <img src=\"https://raw.githubusercontent.com/Athe-kunal/Text-to-SQL/main/Schema.png\" alt=\"Sublime's custom image\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_results(table_collection_,question:str):\n",
    "    # question_emb = emb_fn.embed_with_retries(question)[0]\n",
    "    # Get the table results for the given question\n",
    "    table_results = table_collection_.query(\n",
    "        query_texts = question,\n",
    "        n_results = 5\n",
    "    )\n",
    "    # print(table_results['documents'][0])\n",
    "    return table_results\n",
    "\n",
    "def get_row_results(row_collection_,question,table_name:str):\n",
    "    # Get the row results for the given question\n",
    "    row_results = row_collection_.query(\n",
    "        query_texts = question,\n",
    "        where = {\"table_name\":table_name},\n",
    "        n_results = 5\n",
    "    )\n",
    "    print(row_results['documents'][0])\n",
    "    return row_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: margin \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Gross_Margin (INTEGER) | Net_Margin (INTEGER) | Pre_tax_Pre_stock_compensation_Operating_Margin (INTEGER) | Pre_tax_Unadjusted_Operating_Margin (INTEGER) | After_tax_Unadjusted_Operating_Margin (INTEGER) | Pre_tax_Lease_adjusted_Margin (INTEGER) | After_tax_Lease_Adjusted_Margin (INTEGER) | Pre_tax_Lease_R_D_adj_Margin (INTEGER) | After_tax_Lease_R_D_adj_Margin (INTEGER) | EBITDA_Sales (INTEGER) | EBITDASG_A_Sales (INTEGER) | EBITDAR_D_Sales (INTEGER) | COGS_Sales (INTEGER) | R_D_Sales (INTEGER) | SG_A_Sales (INTEGER) | Stock_Based_Compensation_Sales (INTEGER) | Lease_Expense_Sales (INTEGER)\n",
      "row 1 : Packaging & Container | 22 | 0.2171338247779519 | 0.0285269723092998 | 0.1016364547958028 | 0.0975946017284918 | 0.0799113339135634 | 0.0997355299061771 | 0.0816643450787498 | 0.0998025797146062 | 0.0817313948871789 | 0.1571078511407175 | 0.2522079768591083 | 0.1617111987330198 | 0.782866175222048 | 0.0046033475923023 | 0.0951001257183908 | 0.004041853067311 | 0.0111359826715\n",
      "row 2 : Software (Entertainment) | 84 | 0.6343190737595645 | 0.203484453534525 | 0.3536120749157563 | 0.2680484859792075 | 0.2543465254185649 | 0.2650138705989933 | 0.2514670319003695 | 0.2816776939314963 | 0.2681308552328725 | 0.290337329951565 | 0.4724580252979131 | 0.4795894540185941 | 0.3656809262404354 | 0.1892521240670291 | 0.182120695346348 | 0.0855635889365488 | 0.0140141013588\n",
      "row 3 : Software (Internet) | 35 | 0.591108552762756 | -0.1431970955040302 | 0.1239579652577063 | -0.0335948165704685 | -0.0327194780747064 | -0.0289602901492322 | -0.0282057077640929 | -0.0083175074368089 | -0.0075629250516695 | 0.0122636239274153 | 0.4161715858041718 | 0.2226547217998442 | 0.408891447237244 | 0.2103910978724289 | 0.4039079618767565 | 0.1575527818281749 | 0.02956631949802\n",
      "row 4 : Software (System & Application) | 351 | 0.7152038033402293 | 0.1914080753155209 | 0.3405275162267305 | 0.2529865559329903 | 0.2423866072527709 | 0.2534208835158065 | 0.2428027368326486 | 0.2629596189730837 | 0.2523414722899257 | 0.2853446560160495 | 0.5664609609773408 | 0.4563813852186625 | 0.2847961966597707 | 0.171036729202613 | 0.2811163049612913 | 0.0875409602937401 | 0.01484216535122\n",
      "row 5 : Electronics (Consumer & Office) | 13 | 0.323558031288744 | -0.0304513738086002 | 0.0289465210995791 | -0.0007643401586021 | -0.0007084824743491 | -0.0004205933789403 | -0.0003898565768825 | 0.0056724354829463 | 0.005703172285004 | 0.0359642158837002 | 0.2637004492123042 | 0.130762045590879 | 0.676441968711256 | 0.0947978297071788 | 0.2277362333286039 | 0.0297108612581812 | 0.00871672702736\n",
      "*/\n",
      "\n",
      "Table name: EVA \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | ROE (INTEGER) | Cost_of_Equity (INTEGER) | _ROE_COE_ (INTEGER) | BV_of_Equity (INTEGER) | Equity_EVA_US_millions_ (INTEGER) | ROC (INTEGER) | Cost_of_Capital (INTEGER) | _ROC_WACC_ (INTEGER) | BV_of_Capital (INTEGER) | EVA_US_millions_ (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER)\n",
      "row 1 : Packaging & Container | 22 | 1.1327220379604754 | 0.0854306569307375 | 0.0909052137461818 | -0.0054745568154443 | 49084.58099999999 | -268.71632744677845 | 0.1421364757846558 | 0.0708610226399319 | 0.0712754531447238 | 84525.26117167753 | 6024.57629218745 | 0.620115895714898 | 0.2624270896574677 | 0.050855 | 0.1811910443993945 | 0.03814125 | 0.3798841042851\n",
      "row 2 : Software (Entertainment) | 84 | 1.1083305740238332 | 0.2325593855621183 | 0.0897832064050963 | 0.1427761791570219 | 403869.63 | 57662.96264896017 | 0.1994487340625551 | 0.088266127488506 | 0.1111826065740491 | 620525.0853364643 | 68991.59643229235 | 0.9694385050656324 | 0.5278932207822914 | 0.053524 | 0.0511174704478853 | 0.040143 | 0.03056149493436\n",
      "row 3 : Software (Internet) | 35 | 1.6160843271745926 | -0.1454065112919735 | 0.1131398790500312 | -0.2585463903420048 | 28399.863 | -7342.682064857458 | -0.0048109532624457 | 0.1059203603347457 | -0.1107313135971915 | 45334.05312597519 | -5019.899253324099 | 0.8930134645020447 | 0.65221602599035 | 0.060879 | 0.0260557605345441 | 0.04565925 | 0.10698653549795\n",
      "row 4 : Software (System & Application) | 351 | 1.2939647291542515 | 0.2444183538229951 | 0.0983223775410955 | 0.1460959762818996 | 397835.59000000014 | 58122.178920735554 | 0.2322438986671061 | 0.0949239234411956 | 0.1373199752259105 | 551977.6712097155 | 75797.5601357739 | 0.9415866198035656 | 0.5208648982480728 | 0.053524 | 0.0418992568246478 | 0.040143 | 0.05841338019643\n",
      "row 5 : Electronics (Consumer & Office) | 13 | 1.29517877665579 | -0.0662680013872474 | 0.0983782237261663 | -0.1646462251134137 | 2969.9099999999994 | -488.9844704265786 | 0.009542635546224 | 0.0890493344743714 | -0.0795066989281474 | 3862.6772748588 | -307.1087191487957 | 0.845130180440282 | 0.3942415498208712 | 0.050855 | 0.0730796146510295 | 0.03814125 | 0.15486981955971\n",
      "*/\n",
      "\n",
      "Table name: DollarUS \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Average_Company_Age_years_ (INTEGER) | Market_Cap_millions_ (INTEGER) | Book_Equity_millions_ (INTEGER) | Enteprise_Value_millions_ (INTEGER) | Invested_Capital_millions_ (INTEGER) | Total_Debt_including_leases_millions_ (INTEGER) | Revenues_millions_ (INTEGER) | Gross_Profit_millions_ (INTEGER) | EBITDA_millions_ (INTEGER) | EBIT_Operating_Income_millions_ (INTEGER) | Net_Income_millions_ (INTEGER)\n",
      "row 1 : Packaging & Container | 22 | 80.19047619047619 | 132627.167 | 49084.58099999999 | 205727.3441716776 | 86259.32917167754 | 81247.63917167754 | 146995.2 | 31917.63 | 23820.874 | 14660.644165664493 | 4193.3279999999\n",
      "row 2 : Software (Entertainment) | 84 | 14.544117647058824 | 2773080.8990000025 | 403869.63 | 2781007.513051708 | 383348.20548245066 | 87421.2210517081 | 461576.65299999993 | 292786.8749999999 | 156210.96800000002 | 122324.21538965842 | 93923.673000000\n",
      "row 3 : Software (Internet) | 35 | 14.413793103448276 | 220091.11299999992 | 28399.863 | 241745.7304766904 | 25653.71492597519 | 26367.783476690336 | 28838.050000000003 | 17046.417999999998 | 1980.7540000000004 | -835.1582953380669 | -4129.5\n",
      "row 4 : Software (System & Application) | 351 | 20.125 | 5277974.335000003 | 397835.59000000014 | 5446717.280701609 | 326815.0326293967 | 327430.6527016032 | 508015.76600000006 | 363334.808 | 170251.41100000014 | 128741.80425967928 | 97238.320000000\n",
      "row 5 : Electronics (Consumer & Office) | 13 | 35.61538461538461 | 5288.128000000001 | 2969.9099999999994 | 5468.249666410444 | 2031.147666410444 | 969.0476664104436 | 6463.090999999999 | 2091.185 | 195.229 | -2.7183332820887505 | -196.\n",
      "*/\n",
      "\n",
      "Table name: dbtfund \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Book_Debt_to_Capital (INTEGER) | Market_Debt_to_Capital_Unadjusted_ (INTEGER) | Market_D_E_unadjusted_ (INTEGER) | Market_Debt_to_Capital_adjusted_for_leases_ (INTEGER) | Market_D_E_adjusted_for_leases_ (INTEGER) | Interest_Coverage_Ratio (INTEGER) | Debt_to_EBITDA (INTEGER) | Effective_tax_rate (INTEGER) | Institutional_Holdings (INTEGER) | Std_dev_in_Stock_Prices (INTEGER) | EBITDA_EV (INTEGER) | Net_PP_E_Total_Assets (INTEGER) | Capital_Spending_Total_Assets (INTEGER)\n",
      "row 1 : Packaging & Container | 22 | 0.6103387810613522 | 0.3782286453204903 | 0.6083082510538735 | 0.379884104285102 | 0.6126017844570075 | 3.764310529409295 | 3.5181123824560183 | 0.1811910443993945 | 0.6783750000000001 | 0.2624270896574677 | 0.1157885651803374 | 0.3623114386831347 | 0.04952312081083\n",
      "row 2 : Software (Entertainment) | 84 | 0.1668568483248547 | 0.0310823698127216 | 0.0320794759475208 | 0.0305614949343675 | 0.0315249443617865 | 82.70045339060373 | 0.6523342120398788 | 0.0511174704478853 | 0.380174716981132 | 0.5278932207822914 | 0.0561706386145586 | 0.3664219490280888 | 0.08632671853148\n",
      "row 3 : Software (Internet) | 35 | 0.5188435528953421 | 0.1069721233683309 | 0.1197858724990863 | 0.1069865354979553 | 0.1198039444540877 | -1.2424451113163029 | 74.55708316963612 | 0.0260557605345441 | 0.503754090909091 | 0.65221602599035 | 0.0081935428439386 | 0.1393637519325522 | 0.02491814851918\n",
      "row 4 : Software (System & Application) | 351 | 0.4037655070426813 | 0.0578505240039251 | 0.0614027025199621 | 0.0584133801964343 | 0.0620371816759891 | 12.091079193408472 | 2.258772022287284 | 0.0418992568246478 | 0.4265529149797569 | 0.5208648982480728 | 0.0312576185298293 | 0.1534323819146964 | 0.03979088603498\n",
      "row 5 : Electronics (Consumer & Office) | 13 | 0.257103532061476 | 0.1538670284938941 | 0.1818473380372033 | 0.1548698195597179 | 0.1832496615835402 | -0.0653906229317238 | 4.169022829162121 | 0.0730796146510295 | 0.322611111111111 | 0.3942415498208712 | 0.0357022835294488 | 0.0758014027166977 | 0.01817671397788\n",
      "*/\n",
      "\n",
      "Table name: finflows \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Dividends_in_millions (INTEGER) | Buybacks_in_millions (INTEGER) | Equity_Issuance_in_millions (INTEGER) | Net_Equity_Change_in_millions (INTEGER) | Net_Equity_Change_as_of_Book_Equity (INTEGER) | Debt_Repaid_in_millions (INTEGER) | Debt_Raised_in_millions (INTEGER) | Net_Debt_Change_in_millions (INTEGER) | Net_Change_in_Debt_as_of_Total_Debt (INTEGER) | Change_in_Lease_Debt_in_millions (INTEGER)\n",
      "row 1 : Packaging & Container | 22 | 2944.63 | 2233.409 | 88.02499999999999 | -2145.384 | -0.0437079008579089 | 21249.516000000007 | 23186.13 | 1936.6139999999905 | 0.0252378854855187 | 1078.3980000000\n",
      "row 2 : Software (Entertainment) | 84 | 235.41 | 90230.608 | 492.176 | -89738.43199999999 | -0.2221965340647178 | 19934.036 | 26704.583 | 6770.546999999999 | 0.1247056048135642 | 6193.5280000000\n",
      "row 3 : Software (Internet) | 35 | 0.256 | 5130.625999999999 | 532.7600000000001 | -4597.865999999999 | -0.1618974711251248 | 5242.0070000000005 | 5173.334999999999 | -68.67200000000139 | -0.0029087745986446 | 1487.\n",
      "row 4 : Software (System & Application) | 351 | 26468.845 | 51669.898 | 11749.602 | -39920.296 | -0.1003437022816384 | 35002.797 | 51147.364 | 16144.567000000005 | 0.0548813000879868 | 3688.94599999997\n",
      "row 5 : Electronics (Consumer & Office) | 13 | 0.0 | 193.63 | 36.461000000000006 | -157.16899999999998 | -0.0529204588691239 | 397.182 | 379.297 | -17.885000000000048 | -0.023463092038029 | 74.409999999999\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to find the EBITDA for the software industry, which includes Software (Entertainment), Software (Internet), and Software (System & Application).',\n",
      "    sql=\"```sql\\nSELECT DISTINCT DollarUS.Industry_Name AS Industry, DollarUS.EBITDA_millions_ AS EBITDA\\nFROM DollarUS\\nWHERE DollarUS.Industry_Name LIKE 'Software%'\\n```\"\n",
      ")\n",
      "Extracted rows:  Industry = Software (Entertainment), EBITDA = 156210.96800000002\n",
      " Industry = Software (Internet), EBITDA = 1980.7540000000004\n",
      " Industry = Software (System & Application), EBITDA = 170251.41100000014\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from typing import Any\n",
    "import re\n",
    "\n",
    "class TextToSQLQueryModule(dspy.Module):\n",
    "    \"\"\"Text to SQL to final module\"\"\"\n",
    "    def __init__(self,region:str,use_cot:bool=True,max_retries:int=3):\n",
    "        \"\"\"Text to Answer init module\n",
    "\n",
    "        Args:\n",
    "            region (str): Region for which the module will be used.\n",
    "            use_cot (bool, optional): Whether to use chain of thought for sql query generation. Defaults to True.\n",
    "            max_retries (int, optional): Number of max retries for SQLError. Defaults to 3.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.region = region\n",
    "        db,table_collection,row_collection = db_collection_dict[region]\n",
    "        # print(db,table_collection,row_collection)\n",
    "        self.table_collection = table_collection\n",
    "        self.use_cot = use_cot\n",
    "        self.db = db\n",
    "        self.row_collection = row_collection\n",
    "        if self.use_cot == True:\n",
    "            self.sqlAnswer = dspy.ChainOfThought(TextToSQLAnswer)\n",
    "        else:\n",
    "            self.sqlAnswer = dspy.Predict(TextToSQLAnswer)\n",
    "        self.final_output = dspy.Predict(SQLReturnToAnswer)\n",
    "        self.max_tries = max_retries\n",
    "        # Initialize the sql rectifier with CoT reasoning\n",
    "        self.sql_rectifier = dspy.ChainOfThought(SQLRectifier,rationale_type=dspy.OutputField(\n",
    "            prefix=\"Reasoning: Let's think step by step in order to\",\n",
    "            desc=\"${produce the answer}. We ...\"\n",
    "        ))\n",
    "    \n",
    "    def __call__(self, *args: Any, **kwargs: Any) -> Any:\n",
    "        return self.forward(*args, **kwargs)\n",
    "        \n",
    "    def forward(self,question):\n",
    "        # Embed the question with embedding function\n",
    "        question_emb = emb_fn([question])[0]\n",
    "        # Retrieve the relevant tables from table schema and table summary\n",
    "        docs = self.table_collection.query(\n",
    "            query_embeddings = question_emb,\n",
    "            n_results = 5\n",
    "        )\n",
    "        # docs = get_table_results(db_collection_dict[self.region][1],question)\n",
    "        relevant_rows_schemas = \"\"\n",
    "    \n",
    "        existing_table_names = []\n",
    "\n",
    "        for table_idx,metadata_name in enumerate(docs['metadatas'][0]):\n",
    "            table_metadata = metadata_name['table_metadata']\n",
    "            table_name = metadata_name['table_name']\n",
    "            # If the table name is already in the list of existing table names, skip it\n",
    "            # if table_name in existing_table_names: \n",
    "            #     continue\n",
    "            existing_table_names.append(table_name)\n",
    "            # Retrieve the relevant rows from the current table\n",
    "            rows = self.row_collection.query(\n",
    "                query_embeddings = question_emb,\n",
    "                n_results = 5,\n",
    "                # where clause to filter the rows\n",
    "                where = {\"table_name\":table_name}\n",
    "            )\n",
    "            # Retrieve the relevant table with the schema and summary\n",
    "            relevant_rows_schemas += f'Table name: {table_name} \\n'\n",
    "            relevant_rows_schemas += \"/* \\n\"\n",
    "            for match in re.finditer(\"columns: \",table_metadata):\n",
    "                cols_end = match.end()\n",
    "            relevant_rows_schemas += \"col : \" + \" | \".join(table_metadata[cols_end:].split(\", \")) + \"\\n\"\n",
    "            for row_idx,row in enumerate(rows['metadatas'][0]):\n",
    "                # Get the relevant rows from the current table\n",
    "                # relevant_rows_schemas += f'\\tRow {row_idx+1} from table {table_name}: {row[\"full_rows\"]}\\n'\n",
    "                relevant_rows_schemas += f'row {row_idx+1} : {\" | \".join(row[\"full_rows\"].split(\", \"))}\\n'\n",
    "            relevant_rows_schemas += \"*/\" + '\\n\\n'\n",
    "        print(colored(relevant_rows_schemas,\"yellow\"))\n",
    "        # return \n",
    "        sql_query = self.sqlAnswer(question=question,relevant_table_schemas_rows=relevant_rows_schemas)\n",
    "\n",
    "        num_tries = 0\n",
    "        print(sql_query)\n",
    "        while num_tries <= self.max_tries:\n",
    "            with self.db.connect() as conn:\n",
    "                try:\n",
    "                    # Try executing the sql query for the database\n",
    "                    result = conn.execute(text(process_sql_str(sql_query.sql)))\n",
    "                    num_tries = self.max_tries + 1\n",
    "                except Exception as error:\n",
    "                    # If there is an sql error, then try again with the sql rectifier\n",
    "                    print(colored(str(error),'red'))\n",
    "                    sql_query = self.sql_rectifier(input_sql=sql_query.sql,error_str=str(error),relevant_table_schemas_rows=relevant_rows_schemas)\n",
    "                    print(colored(sql_query.rationale,'green'))\n",
    "                    print()\n",
    "                    print(colored(sql_query.sql,'green'))\n",
    "                    # If all the num_retries are exhausted, then exit the program\n",
    "                    num_tries += 1\n",
    "                    if num_tries == self.max_tries+1:\n",
    "                        return sql_query,error\n",
    "        # With the retrieved rows from the database, then try to answer the question with dspy context\n",
    "        with dspy.context(lm=sql_to_answer):\n",
    "            row_str = \"\"\n",
    "            key = tuple(result.keys())\n",
    "            for row in result.fetchall():\n",
    "                for r,k in zip(row,key):\n",
    "                    row_str += f\" {k} = {r},\"\n",
    "                row_str = row_str[:-1]\n",
    "                row_str += \"\\n\"\n",
    "            print(f\"Extracted rows: {row_str}\")\n",
    "            final_answer = self.final_output(question=question,sql=sql_query.sql,relevant_rows=row_str)\n",
    "            return final_answer\n",
    "tsql_ = TextToSQLQueryModule(\"US\")\n",
    "question = \"What is the ebitda of software and packaging industry?\"\n",
    "sq = tsql_(question = question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction(\n",
      "    answer='The EBITDA of the software industry is as follows:\\n- Software (Entertainment): $156,210.97 million\\n- Software (Internet): $1,980.75 million\\n- Software (System & Application): $170,251.41 million'\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(sq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: EVA \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | ROE (INTEGER) | Cost_of_Equity (INTEGER) | _ROE_COE_ (INTEGER) | BV_of_Equity (INTEGER) | Equity_EVA_US_millions_ (INTEGER) | ROC (INTEGER) | Cost_of_Capital (INTEGER) | _ROC_WACC_ (INTEGER) | BV_of_Capital (INTEGER) | EVA_US_millions_ (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER)\n",
      "row 1 : Aerospace/Defense | 70 | 1.0764050153356324 | 0.1319148365195281 | 0.088314630705439 | 0.043600205814089 | 145775.301 | 6355.833126210783 | 0.1616771766434065 | 0.0781335598681736 | 0.0835436167752329 | 195266.1836785369 | 16313.243218401924 | 0.7970822237983712 | 0.3640214967123067 | 0.050855 | 0.0727529390842114 | 0.03814125 | 0.20291777620162\n",
      "row 2 : Total Market (without financials) | 5214 | 1.097522642247185 | 0.1659255247192018 | 0.0892860415433705 | 0.0766394831758313 | 8566043.05699999 | 656497.1127503973 | 0.1466884216172786 | 0.0798580827010761 | 0.0668303389162024 | 14577190.305961732 | 974198.568593403 | 0.815661408370401 | 0.4700345950094993 | 0.050855 | 0.0680489965490159 | 0.03814125 | 0.1843385916295\n",
      "row 3 : Software (Entertainment) | 84 | 1.1083305740238332 | 0.2325593855621183 | 0.0897832064050963 | 0.1427761791570219 | 403869.63 | 57662.96264896017 | 0.1994487340625551 | 0.088266127488506 | 0.1111826065740491 | 620525.0853364643 | 68991.59643229235 | 0.9694385050656324 | 0.5278932207822914 | 0.053524 | 0.0511174704478853 | 0.040143 | 0.03056149493436\n",
      "row 4 : Electronics (General) | 129 | 0.9339545990672654 | 0.0879766010692752 | 0.0817619115570942 | 0.006214689512181 | 72980.08699999998 | 453.548581276959 | 0.1511438395305996 | 0.0753610851527285 | 0.0757827543778711 | 81592.19108678924 | 6183.280976282475 | 0.8532615926517355 | 0.4294264362055603 | 0.050855 | 0.0817457737542162 | 0.03814125 | 0.14673840734826\n",
      "row 5 : Semiconductor Equip | 30 | 1.5279822074255314 | 0.3289066823779071 | 0.1090871815415744 | 0.2198195008363326 | 44503.16999999999 | 9782.664615034451 | 0.2568363246651193 | 0.1039673141944747 | 0.1528690104706445 | 70041.38797094872 | 10707.15767110944 | 0.9278342360745604 | 0.3704014298865129 | 0.050855 | 0.1214195116015062 | 0.03814125 | 0.07216576392543\n",
      "*/\n",
      "\n",
      "Table name: margin \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Gross_Margin (INTEGER) | Net_Margin (INTEGER) | Pre_tax_Pre_stock_compensation_Operating_Margin (INTEGER) | Pre_tax_Unadjusted_Operating_Margin (INTEGER) | After_tax_Unadjusted_Operating_Margin (INTEGER) | Pre_tax_Lease_adjusted_Margin (INTEGER) | After_tax_Lease_Adjusted_Margin (INTEGER) | Pre_tax_Lease_R_D_adj_Margin (INTEGER) | After_tax_Lease_R_D_adj_Margin (INTEGER) | EBITDA_Sales (INTEGER) | EBITDASG_A_Sales (INTEGER) | EBITDAR_D_Sales (INTEGER) | COGS_Sales (INTEGER) | R_D_Sales (INTEGER) | SG_A_Sales (INTEGER) | Stock_Based_Compensation_Sales (INTEGER) | Lease_Expense_Sales (INTEGER)\n",
      "row 1 : Aerospace/Defense | 70 | 0.1727498953821253 | 0.0496288894345249 | 0.0970899067124733 | 0.0853938467891159 | 0.0791811934555008 | 0.085557881411237 | 0.0793332940767511 | 0.0877011542663323 | 0.0814765669318464 | 0.1213211709923032 | 0.1854721753152544 | 0.163301436750081 | 0.8272501046178746 | 0.0419802657577777 | 0.0641510043229512 | 0.0116960599233574 | 0.00737353454744\n",
      "row 2 : Total Market (without financials) | 5214 | 0.3341124303745417 | 0.0759000807387551 | 0.1349854587183475 | 0.1203243079012881 | 0.1121363594881506 | 0.1193029155861628 | 0.1111844718951524 | 0.122305912015432 | 0.1141874683244217 | 0.1649863910103977 | 0.308807292459029 | 0.2037377626528898 | 0.6658875696254583 | 0.038751371642492 | 0.1438209014486313 | 0.0146611508170593 | 0.01534583398930\n",
      "row 3 : Software (Entertainment) | 84 | 0.6343190737595645 | 0.203484453534525 | 0.3536120749157563 | 0.2680484859792075 | 0.2543465254185649 | 0.2650138705989933 | 0.2514670319003695 | 0.2816776939314963 | 0.2681308552328725 | 0.290337329951565 | 0.4724580252979131 | 0.4795894540185941 | 0.3656809262404354 | 0.1892521240670291 | 0.182120695346348 | 0.0855635889365488 | 0.0140141013588\n",
      "row 4 : Electronics (General) | 129 | 0.2736441477972378 | 0.0469265820090722 | 0.1081365384368007 | 0.0948166055844074 | 0.0870657487961617 | 0.095866096633002 | 0.0880294483869408 | 0.0979701850709138 | 0.0901335368248526 | 0.1483155870270167 | 0.274700126354909 | 0.1987893448489177 | 0.7263558522027622 | 0.0504737578219009 | 0.1263845393278922 | 0.0133199328523933 | 0.00864521049991\n",
      "row 5 : Semiconductor Equip | 30 | 0.4427439580396038 | 0.1793844069139284 | 0.2615904123507109 | 0.2418675138617241 | 0.212500078456363 | 0.243380403940693 | 0.2138292741608367 | 0.2500123642381983 | 0.220461234458342 | 0.3032502002655829 | 0.3930952875683865 | 0.4064476540724488 | 0.5572560419603962 | 0.1031974538068659 | 0.0898450873028035 | 0.0197228984889867 | 0.00906011117075\n",
      "*/\n",
      "\n",
      "Table name: DollarUS \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Average_Company_Age_years_ (INTEGER) | Market_Cap_millions_ (INTEGER) | Book_Equity_millions_ (INTEGER) | Enteprise_Value_millions_ (INTEGER) | Invested_Capital_millions_ (INTEGER) | Total_Debt_including_leases_millions_ (INTEGER) | Revenues_millions_ (INTEGER) | Gross_Profit_millions_ (INTEGER) | EBITDA_millions_ (INTEGER) | EBIT_Operating_Income_millions_ (INTEGER) | Net_Income_millions_ (INTEGER)\n",
      "row 1 : Aerospace/Defense | 70 | 53.82539682539682 | 788306.4130000002 | 145775.301 | 960283.3679021292 | 164231.158307403 | 200683.66790212895 | 387474.417 | 66936.16500000001 | 48918.45400000002 | 33151.49021957421 | 19229.925000000\n",
      "row 2 : Total Market (without financials) | 5214 | 34.79463487696991 | 43163609.00299983 | 8566043.05699999 | 51078189.49081551 | 13583012.82046062 | 9754928.713815568 | 18726267.155 | 6256678.631000005 | 3370108.721999998 | 2234098.269636899 | 1421325.1889999\n",
      "row 3 : Software (Entertainment) | 84 | 14.544117647058824 | 2773080.8990000025 | 403869.63 | 2781007.513051708 | 383348.20548245066 | 87421.2210517081 | 461576.65299999993 | 292786.8749999999 | 156210.96800000002 | 122324.21538965842 | 93923.673000000\n",
      "row 4 : Electronics (General) | 129 | 39.15573770491803 | 269499.1279999999 | 72980.08699999998 | 297860.8621360993 | 66341.67523831947 | 46346.71613609941 | 136820.96000000014 | 37440.255 | 20098.71 | 13116.491372780118 | 6420.5399999999\n",
      "row 5 : Semiconductor Equip | 30 | 42.5 | 414640.35 | 44503.16999999999 | 427920.4277709488 | 48010.57777094871 | 32250.19777094872 | 81597.89499999999 | 36126.97499999999 | 23479.94899999999 | 19859.32864581025 | 14637.\n",
      "*/\n",
      "\n",
      "Table name: wacc \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | Cost_of_Equity (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER) | Cost_of_Capital (INTEGER) | Cost_of_Capital_Local_Currency_ (INTEGER)\n",
      "row 1 : Aerospace/Defense | 70 | 1.0764050153356324 | 0.088314630705439 | 0.7970822237983712 | 0.3640214967123067 | 0.050855 | 0.0727529390842114 | 0.03814125 | 0.2029177762016287 | 0.0781335598681736 | 0.07813355986817\n",
      "row 2 : Total Market (without financials) | 5214 | 1.097522642247185 | 0.0892860415433705 | 0.815661408370401 | 0.4700345950094993 | 0.050855 | 0.0680489965490159 | 0.03814125 | 0.1843385916295989 | 0.0798580827010761 | 0.07985808270107\n",
      "row 3 : Software (Entertainment) | 84 | 1.1083305740238332 | 0.0897832064050963 | 0.9694385050656324 | 0.5278932207822914 | 0.053524 | 0.0511174704478853 | 0.040143 | 0.0305614949343675 | 0.088266127488506 | 0.08826612748850\n",
      "row 4 : Electronics (General) | 129 | 0.9339545990672654 | 0.0817619115570942 | 0.8532615926517355 | 0.4294264362055603 | 0.050855 | 0.0817457737542162 | 0.03814125 | 0.1467384073482645 | 0.0753610851527285 | 0.07536108515272\n",
      "row 5 : Semiconductor Equip | 30 | 1.5279822074255314 | 0.1090871815415744 | 0.9278342360745604 | 0.3704014298865129 | 0.050855 | 0.1214195116015062 | 0.03814125 | 0.0721657639254397 | 0.1039673141944747 | 0.10396731419447\n",
      "*/\n",
      "\n",
      "Table name: dbtfund \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Book_Debt_to_Capital (INTEGER) | Market_Debt_to_Capital_Unadjusted_ (INTEGER) | Market_D_E_unadjusted_ (INTEGER) | Market_Debt_to_Capital_adjusted_for_leases_ (INTEGER) | Market_D_E_adjusted_for_leases_ (INTEGER) | Interest_Coverage_Ratio (INTEGER) | Debt_to_EBITDA (INTEGER) | Effective_tax_rate (INTEGER) | Institutional_Holdings (INTEGER) | Std_dev_in_Stock_Prices (INTEGER) | EBITDA_EV (INTEGER) | Net_PP_E_Total_Assets (INTEGER) | Capital_Spending_Total_Assets (INTEGER)\n",
      "row 1 : Aerospace/Defense | 70 | 0.5598223230984364 | 0.2003176164886141 | 0.250496472873423 | 0.2029177762016287 | 0.2545757139516368 | 3.740440705503212 | 4.269061419331231 | 0.0727529390842114 | 0.485818947368421 | 0.3640214967123067 | 0.0509416862096332 | 0.1258713546270282 | 0.01798339388026\n",
      "row 2 : Total Market (without financials) | 5214 | 0.5001187576177913 | 0.1827761005242874 | 0.2236548645255564 | 0.184338591629599 | 0.225998912953215 | 6.428195471118638 | 3.157364796897844 | 0.0680489965490159 | 0.4702478911478115 | 0.4700345950094993 | 0.0659794083462176 | 0.32642976153906 | 0.04335448376240\n",
      "row 3 : Software (Entertainment) | 84 | 0.1668568483248547 | 0.0310823698127216 | 0.0320794759475208 | 0.0305614949343675 | 0.0315249443617865 | 82.70045339060373 | 0.6523342120398788 | 0.0511174704478853 | 0.380174716981132 | 0.5278932207822914 | 0.0561706386145586 | 0.3664219490280888 | 0.08632671853148\n",
      "row 4 : Electronics (General) | 129 | 0.3587835842228601 | 0.1457673736295513 | 0.1706413090880205 | 0.1467384073482645 | 0.1719735291169455 | 6.641668616399566 | 2.2839129110687453 | 0.0817457737542162 | 0.4456912621359226 | 0.4294264362055603 | 0.0674768408842396 | 0.182597141990381 | 0.03108461726876\n",
      "row 5 : Semiconductor Equip | 30 | 0.3856389110464103 | 0.0719131674240597 | 0.0774853870348122 | 0.0721657639254397 | 0.0777787250347167 | 12.620099818076602 | 1.3033238138451468 | 0.1214195116015062 | 0.6812737931034483 | 0.3704014298865129 | 0.0548698951398693 | 0.1426826921482972 | 0.03734252665127\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the sql. We need to first join the tables based on the Industry_Name column and then filter out the rows for Software industries, semiconductor industry, and aerospace. Finally, we need to select the EBITDA value and count the number of firms.',\n",
      "    sql=\"```sql\\nSELECT DISTINCT DollarUS.Industry_Name AS Industry, DollarUS.EBITDA_millions_ AS EBITDA, DollarUS.Number_of_firms AS Number_of_Firms\\nFROM DollarUS\\nJOIN EVA ON DollarUS.Industry_Name = EVA.Industry_Name\\nJOIN margin ON DollarUS.Industry_Name = margin.Industry_Name\\nWHERE DollarUS.Industry_Name IN ('Software (Entertainment)', 'Semiconductor Equip', 'Aerospace/Defense');\\n```\"\n",
      ")\n",
      "Extracted rows:  Industry = Aerospace/Defense, EBITDA = 48918.45400000002, Number_of_Firms = 70\n",
      " Industry = Semiconductor Equip, EBITDA = 23479.94899999999, Number_of_Firms = 30\n",
      " Industry = Software (Entertainment), EBITDA = 156210.96800000002, Number_of_Firms = 84\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# tsql = TextToSQLQueryModule(\"US\")\n",
    "# sq = tsql(question=\"What is the effective tax rate of the healthcare industry?\")\n",
    "sq = tsql_(question=\"What is the EBITDA value and number of firms for all the Software industries, semiconductor industry and aerospace?\")\n",
    "# sq = tsql(\"What is the debt to EBITDA ratio for software industry?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The EBITDA value for the Software (Entertainment) industry is $156,210.97 million with 84 firms. For the Semiconductor Equip industry, the EBITDA value is $23,479.95 million with 30 firms. Lastly, for the Aerospace/Defense industry, the EBITDA value is $48,918.45 million with 70 firms.\n"
     ]
    }
   ],
   "source": [
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: betaIndia \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Beta_ (INTEGER) | D_E_Ratio (INTEGER) | Effective_Tax_rate (INTEGER) | Unlevered_beta (INTEGER) | Cash_Firm_value (INTEGER) | Unlevered_beta_corrected_for_cash (INTEGER) | HiLo_Risk (INTEGER) | Standard_deviation_of_equity (INTEGER) | Standard_deviation_in_operating_income_last_10_years_ (INTEGER) | Beta_2020 (INTEGER) | Beta_2021 (INTEGER) | Beta_2022 (INTEGER) | Beta_2023 (INTEGER) | Average_Beta_2019_23 (INTEGER)\n",
      "row 1 : Semiconductor | 8 | 1.611640106554315 | 0.0408116499112951 | 0.1361426723519776 | 1.566877312715955 | 0.0034062258080553 | 1.572232692288697 | 0.5628222021115665 | 0.5010518370389238 | 12.74567091312641 | 0.73 | 1.16 | 1.9836144211648223 | 1.9836144211648223 | 1.48589230692366\n",
      "row 2 : Total Market (without financials) | 3850 | 0.8198477825445601 | 0.1608082521880037 | 0.1670198672216159 | 0.7368982579173127 | 0.0296361733703686 | 0.7594040891618811 | 0.384976556985476 | 0.3699643172605134 | 0.4325539868550488 | 0.62 | 0.88 | 0.8646047974287592 | 0.8646047974287592 | 0.79772273680387\n",
      "row 3 : Software (Entertainment) | 5 | 0.7830360024669405 | 0.0075317267581366 | 0.2752644087938206 | 0.7789293244818127 | 0.0099775560297291 | 0.7867794606333218 | 0.2025157685180063 | 0.1867409361734939 | 0.698956301092383 | 0.59 | 1.16 | 1.219451180957866 | 1.219451180957866 | 0.99513636450981\n",
      "row 4 : Software (Internet) | 6 | -0.0057433681694908 | 0.0336472713198297 | 0.0983609410926519 | -0.0056112069085913 | 0.0328328215881574 | -0.0058016928550091 | 0.512617043205324 | 0.504931470076568 | 10.948979663023987 | 1.04 | 0.52 | 0.4883816218840528 | 0.4883816218840528 | 0.50619231018261\n",
      "row 5 : Software (System & Application) | 73 | 1.164011103524074 | 0.014057190040673 | 0.1609424721087531 | 1.1526688044177154 | 0.0427366294396399 | 1.2041292290782728 | 0.3973005076855119 | 0.4048728665269909 | 0.3296550990654787 | 0.83 | 1.2 | 1.2407987165306618 | 1.2407987165306618 | 1.14314533242791\n",
      "*/\n",
      "\n",
      "Table name: totalbetaIndia \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Average_Unlevered_Beta (INTEGER) | Average_Levered_Beta (INTEGER) | Average_correlation_with_the_market (INTEGER) | Total_Unlevered_Beta (INTEGER) | Total_Levered_Beta (INTEGER)\n",
      "row 1 : Semiconductor | 8 | 1.572232692288697 | 1.611640106554315 | 0.1520316120175535 | 10.34148537547025 | 10.6006907719181\n",
      "row 2 : Total Market (without financials) | 3850 | 0.7594040891618811 | 0.8198477825445601 | 0.138919459674997 | 5.466506211142138 | 5.9016050340434\n",
      "row 3 : Software (Entertainment) | 5 | 0.7867794606333218 | 0.7830360024669405 | 0.2362678264192356 | 3.3300321612018964 | 3.3141880311605\n",
      "row 4 : Software (Internet) | 6 | -0.0058016928550091 | -0.0057433681694908 | 0.0877398209734128 | -0.0661238282759565 | -0.06545908238439\n",
      "row 5 : Software (System & Application) | 73 | 1.2041292290782728 | 1.164011103524074 | 0.1534528684231733 | 7.846899451613207 | 7.5854633118626\n",
      "*/\n",
      "\n",
      "Table name: psIndia \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Price_Sales (INTEGER) | Net_Margin (INTEGER) | EV_Sales (INTEGER) | Pre_tax_Operating_Margin (INTEGER)\n",
      "row 1 : Semiconductor | 8 | 6.380530138666162 | 0.0160833883595887 | 6.618309593434581 | 0.04773134609942\n",
      "row 2 : Total Market (without financials) | 3850 | 2.210898698094238 | 0.0681170556023806 | 2.490370305272766 | 0.11262992702128\n",
      "row 3 : Semiconductor Equip | 1 | 4.350600126342388 | 0.0243840808591282 | 4.593177511054959 | 0.03006948831332\n",
      "row 4 : Software (Entertainment) | 5 | 17.543944376879494 | 0.078800724226794 | 17.499716487891103 | 0.20023775912029\n",
      "row 5 : Software (Internet) | 6 | 11.275821937736398 | 0.1813354669770148 | 11.272548734361362 | 0.23581611870817\n",
      "*/\n",
      "\n",
      "Table name: waccIndia \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | Cost_of_Equity (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER) | Cost_of_Capital (INTEGER) | Cost_of_Capital_Local_Currency_ (INTEGER)\n",
      "row 1 : Semiconductor | 8 | 1.611640106554315 | 0.164669092321892 | 0.9607886307625656 | 0.5010518370389238 | 0.077424 | 0.1361426723519776 | 0.0541968 | 0.0392113692374344 | 0.1603373224771525 | 0.18286814427282\n",
      "row 2 : Total Market (without financials) | 3850 | 0.8198477825445601 | 0.1028301118167301 | 0.8614687207082679 | 0.3699643172605134 | 0.074755 | 0.1670198672216159 | 0.0523285 | 0.1385312792917321 | 0.095834058925464 | 0.11711239016673\n",
      "row 3 : Software (Entertainment) | 5 | 0.7830360024669405 | 0.099955111792668 | 0.9925245760921386 | 0.1867409361734939 | 0.068943 | 0.2752644087938206 | 0.0482601 | 0.0074754239078613 | 0.0995686696655959 | 0.12091951762026\n",
      "row 4 : Software (Internet) | 6 | -0.0057433681694908 | 0.0383514429459627 | 0.967448014179086 | 0.504931470076568 | 0.077424 | 0.0983609410926519 | 0.0541968 | 0.0325519858209141 | 0.0388672407841131 | 0.05903942021681\n",
      "row 5 : Software (System & Application) | 73 | 1.164011103524074 | 0.1297092671852301 | 0.9861376752921508 | 0.4048728665269909 | 0.074755 | 0.1609424721087531 | 0.0523285 | 0.0138623247078493 | 0.128636589864366 | 0.15055186345396\n",
      "*/\n",
      "\n",
      "Table name: pbvIndia \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | PBV (INTEGER) | ROE (INTEGER) | EV_Invested_Capital (INTEGER) | ROIC (INTEGER)\n",
      "row 1 : Semiconductor | 8 | 8.627551020408164 | 0.0218198105963654 | 7.1386259299656505 | 0.04466875019704\n",
      "row 2 : Total Market (without financials) | 3850 | 4.152459768346949 | 0.1526067560105244 | 3.119937106367408 | 0.12893652513884\n",
      "row 3 : Semiconductor Equip | 1 | 5.957612456747405 | 0.0373307543520309 | 4.721005520284547 | 0.02701816743642\n",
      "row 4 : Software (Entertainment) | 5 | 2.9646582024861243 | 0.0182539574699368 | 3.0371036678952725 | 0.03630006976547\n",
      "row 5 : Software (Internet) | 6 | 10.946194040389774 | 0.2266569688153468 | 10.876201838725525 | 0.24757464461570\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to retrieve the beta value and number of firms for the Software industries and semiconductor industry. We have the relevant information in the tables betaIndia and totalbetaIndia.',\n",
      "    sql=\"```sql\\nSELECT DISTINCT betaIndia.Industry_Name AS Industry, betaIndia.Beta_ AS Beta_Value, betaIndia.Number_of_firms AS Number_of_Firms\\nFROM betaIndia\\nWHERE betaIndia.Industry_Name LIKE 'Software%' OR betaIndia.Industry_Name = 'Semiconductor'\\n\\nUNION\\n\\nSELECT DISTINCT totalbetaIndia.Industry_Name AS Industry, totalbetaIndia.Average_Unlevered_Beta AS Beta_Value, totalbetaIndia.Number_of_firms AS Number_of_Firms\\nFROM totalbetaIndia\\nWHERE totalbetaIndia.Industry_Name LIKE 'Software%' OR totalbetaIndia.Industry_Name = 'Semiconductor';\\n```\"\n",
      ")\n",
      "Extracted rows:  Industry = Semiconductor, Beta_Value = 1.572232692288697, Number_of_Firms = 8\n",
      " Industry = Semiconductor, Beta_Value = 1.611640106554315, Number_of_Firms = 8\n",
      " Industry = Software (Entertainment), Beta_Value = 0.7830360024669405, Number_of_Firms = 5\n",
      " Industry = Software (Entertainment), Beta_Value = 0.7867794606333218, Number_of_Firms = 5\n",
      " Industry = Software (Internet), Beta_Value = -0.0058016928550091, Number_of_Firms = 6\n",
      " Industry = Software (Internet), Beta_Value = -0.0057433681694908, Number_of_Firms = 6\n",
      " Industry = Software (System & Application), Beta_Value = 1.164011103524074, Number_of_Firms = 73\n",
      " Industry = Software (System & Application), Beta_Value = 1.2041292290782728, Number_of_Firms = 73\n",
      "\n",
      "The beta value and number of firms for the Software industries and semiconductor industry are as follows:\n",
      "\n",
      "- Industry: Semiconductor, Beta Value: 1.572232692288697, Number of Firms: 8\n",
      "- Industry: Semiconductor, Beta Value: 1.611640106554315, Number of Firms: 8\n",
      "- Industry: Software (Entertainment), Beta Value: 0.7830360024669405, Number of Firms: 5\n",
      "- Industry: Software (Entertainment), Beta Value: 0.7867794606333218, Number of Firms: 5\n",
      "- Industry: Software (Internet), Beta Value: -0.0058016928550091, Number of Firms: 6\n",
      "- Industry: Software (Internet), Beta Value: -0.0057433681694908, Number of Firms: 6\n",
      "- Industry: Software (System & Application), Beta Value: 1.164011103524074, Number of Firms: 73\n",
      "- Industry: Software (System & Application), Beta Value: 1.2041292290782728, Number of Firms: 73\n"
     ]
    }
   ],
   "source": [
    "tsql = TextToSQLQueryModule(\"India\")\n",
    "sq = tsql(question=\"What is the beta value and number of firms for all the Software industries and semiconductor industry?\")\n",
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The beta value and number of firms for the Software industries and semiconductor industry are as follows:\n",
      "\n",
      "- Industry: Semiconductor, Beta Value: 1.572232692288697, Number of Firms: 8\n",
      "- Industry: Semiconductor, Beta Value: 1.611640106554315, Number of Firms: 8\n",
      "- Industry: Software (Entertainment), Beta Value: 0.7830360024669405, Number of Firms: 5\n",
      "- Industry: Software (Entertainment), Beta Value: 0.7867794606333218, Number of Firms: 5\n",
      "- Industry: Software (Internet), Beta Value: -0.0058016928550091, Number of Firms: 6\n",
      "- Industry: Software (Internet), Beta Value: -0.0057433681694908, Number of Firms: 6\n",
      "- Industry: Software (System & Application), Beta Value: 1.164011103524074, Number of Firms: 73\n",
      "- Industry: Software (System & Application), Beta Value: 1.2041292290782728, Number of Firms: 73\n"
     ]
    }
   ],
   "source": [
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: betaChina \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Beta_ (INTEGER) | D_E_Ratio (INTEGER) | Effective_Tax_rate (INTEGER) | Unlevered_beta (INTEGER) | Cash_Firm_value (INTEGER) | Unlevered_beta_corrected_for_cash (INTEGER) | HiLo_Risk (INTEGER) | Standard_deviation_of_equity (INTEGER) | Standard_deviation_in_operating_income_last_10_years_ (INTEGER)\n",
      "row 1 : Semiconductor | 173 | 1.3766066410015918 | 0.1214640196981001 | 0.0455831075871332 | 1.261670924483595 | 0.1353894923931989 | 1.4592361686371793 | 0.3256484142873909 | 0.3505300755906819 | 0.93670130463571\n",
      "row 2 : Total Market (without financials) | 7161 | 1.0686360714269856 | 0.4705306537504954 | 0.0996397852701986 | 0.7898866574410776 | 0.1404090933852625 | 0.9189099737592956 | 0.3186107366027715 | 0.3211951486316037 | 0.26343082981471\n",
      "row 3 : Semiconductor Equip | 59 | 1.1458070496329555 | 0.1005963129277871 | 0.072001108546688 | 1.0654237722383604 | 0.0896726175453628 | 1.170374299151055 | 0.3370962305052479 | 0.3780387083633861 | 1.05671095308164\n",
      "row 4 : Software (Entertainment) | 33 | 1.9284684177392413 | 0.134930367400854 | 0.0960440405122507 | 1.7512461980545997 | 0.0533537078374321 | 1.849947770939832 | 0.5242121688667037 | 0.4696202053064469 | 0.48614502813710\n",
      "row 5 : Software (Internet) | 19 | 0.9517387989546392 | 0.2511530032375074 | 0.0593218951020171 | 0.8008810401099454 | 0.0847676071125495 | 0.8750575769977502 | 0.3425492316782614 | 0.3399832570802881 | 0.42653552101587\n",
      "*/\n",
      "\n",
      "Table name: totalbetaChina \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Average_Unlevered_Beta (INTEGER) | Average_Levered_Beta (INTEGER) | Average_correlation_with_the_market (INTEGER) | Total_Unlevered_Beta (INTEGER) | Total_Levered_Beta (INTEGER)\n",
      "row 1 : Semiconductor | 173 | 1.4592361686371793 | 1.3766066410015918 | 0.2051159812465771 | 7.114200267423239 | 6.7113573142149\n",
      "row 2 : Total Market (without financials) | 7161 | 0.9189099737592956 | 1.0686360714269856 | 0.1807939202400685 | 5.082637582829746 | 5.9107965024929\n",
      "row 3 : Semiconductor Equip | 59 | 1.170374299151055 | 1.1458070496329555 | 0.1422791350917505 | 8.225902542887432 | 8.053233166578\n",
      "row 4 : Software (Entertainment) | 33 | 1.849947770939832 | 1.9284684177392413 | 0.189683512387109 | 9.75281271239026 | 10.1667688112164\n",
      "row 5 : Software (Internet) | 19 | 0.8750575769977502 | 0.9517387989546392 | 0.1824645387000501 | 4.795767896776043 | 5.2160206346679\n",
      "*/\n",
      "\n",
      "Table name: EVAChina \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | ROE (INTEGER) | Cost_of_Equity (INTEGER) | _ROE_COE_ (INTEGER) | BV_of_Equity (INTEGER) | Equity_EVA (INTEGER) | ROC (INTEGER) | Cost_of_Capital (INTEGER) | _ROC_WACC_ (INTEGER) | BV_of_Capital (INTEGER) | EVA (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER)\n",
      "row 1 : Semiconductor | 173 | 1.3766066410015918 | 0.0962582854914372 | 0.1163029538883896 | -0.0200446683969523 | 126261.36999999998 | -2530.8672929949053 | 0.0928822041397873 | 0.1084628628354935 | -0.0155806586957062 | 138747.35577016464 | -2161.775195186669 | 0.8916915589224178 | 0.3505300755906819 | 0.0585549999999999 | 0.0455831075871332 | 0.04391625 | 0.10830844107758\n",
      "row 2 : Total Market (without financials) | 7161 | 1.0686360714269856 | 0.0639810392810329 | 0.0989642108213393 | -0.0349831715403063 | 7039463.864000044 | -246262.7719061013 | 0.0650803131143393 | 0.0813503291067108 | -0.0162700159923715 | 10535415.2412372 | -171411.37446120442 | 0.6800266267483532 | 0.3211951486316037 | 0.0585549999999999 | 0.0996397852701986 | 0.04391625 | 0.31997337325164\n",
      "row 3 : Semiconductor Equip | 59 | 1.1458070496329555 | 0.2052549059382613 | 0.1033089368943354 | 0.1019459690439259 | 47325.30000000001 | 4824.623568794508 | 0.1911922171025874 | 0.0978803476411596 | 0.0933118694614277 | 50881.7695261246 | 4747.873035988193 | 0.9085983555040426 | 0.3780387083633861 | 0.0585549999999999 | 0.072001108546688 | 0.04391625 | 0.09140164449595\n",
      "row 4 : Software (Entertainment) | 33 | 1.9284684177392413 | 0.2294528801199007 | 0.1473727719187193 | 0.0820801082011814 | 117637.29 | 9655.68149169376 | 0.0906573095366518 | 0.1350729631256245 | -0.0444156535889726 | 176984.04958861455 | -7860.862237301461 | 0.8811113251733117 | 0.4696202053064469 | 0.0585549999999999 | 0.0960440405122507 | 0.04391625 | 0.11888867482668\n",
      "row 5 : Software (Internet) | 19 | 0.9517387989546392 | 0.0250537522727571 | 0.0923828943811461 | -0.067329142108389 | 8376.390000000001 | -563.9751526652892 | 0.034746695140248 | 0.0826538338572364 | -0.0479071387169884 | 11768.570206159107 | -563.7985253670817 | 0.7992627579619607 | 0.3399832570802881 | 0.0585549999999999 | 0.0593218951020171 | 0.04391625 | 0.20073724203803\n",
      "*/\n",
      "\n",
      "Table name: waccChina \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | Cost_of_Equity (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER) | Cost_of_Capital (INTEGER) | Cost_of_Capital_Local_Currency_ (INTEGER)\n",
      "row 1 : Semiconductor | 173 | 1.3766066410015918 | 0.1163029538883896 | 0.8916915589224178 | 0.3505300755906819 | 0.0585549999999999 | 0.0455831075871332 | 0.04391625 | 0.1083084410775822 | 0.1084628628354936 | 0.10846286283549\n",
      "row 2 : Total Market (without financials) | 7161 | 1.0686360714269856 | 0.0989642108213393 | 0.6800266267483532 | 0.3211951486316037 | 0.0585549999999999 | 0.0996397852701986 | 0.04391625 | 0.3199733732516467 | 0.0813503291067108 | 0.08135032910671\n",
      "row 3 : Semiconductor Equip | 59 | 1.1458070496329555 | 0.1033089368943354 | 0.9085983555040426 | 0.3780387083633861 | 0.0585549999999999 | 0.072001108546688 | 0.04391625 | 0.0914016444959574 | 0.0978803476411596 | 0.09788034764115\n",
      "row 4 : Software (Entertainment) | 33 | 1.9284684177392413 | 0.1473727719187193 | 0.8811113251733117 | 0.4696202053064469 | 0.0585549999999999 | 0.0960440405122507 | 0.04391625 | 0.1188886748266883 | 0.1350729631256245 | 0.13507296312562\n",
      "row 5 : Software (Internet) | 19 | 0.9517387989546392 | 0.0923828943811461 | 0.7992627579619607 | 0.3399832570802881 | 0.0585549999999999 | 0.0593218951020171 | 0.04391625 | 0.2007372420380393 | 0.0826538338572364 | 0.08265383385723\n",
      "*/\n",
      "\n",
      "Table name: pbvChina \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | PBV (INTEGER) | ROE (INTEGER) | EV_Invested_Capital (INTEGER) | ROIC (INTEGER)\n",
      "row 1 : Semiconductor | 173 | 2.73256217631175 | 0.0962582854914372 | 2.4500098154005925 | 0.09288220413978\n",
      "row 2 : Total Market (without financials) | 7161 | 1.4084499744108725 | 0.0639810392810329 | 1.2430260813341163 | 0.06508031311433\n",
      "row 3 : Semiconductor Equip | 59 | 2.954264696454902 | 0.2052549059382613 | 2.6326043462726965 | 0.19119221710258\n",
      "row 4 : Software (Entertainment) | 33 | 3.1135265373627137 | 0.2294528801199007 | 2.248186099053892 | 0.09065730953665\n",
      "row 5 : Software (Internet) | 19 | 2.6160739269602717 | 0.0250537522727571 | 2.044727771310165 | 0.0347466951402\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to retrieve the beta value and number of firms for the Software industries and semiconductor industry. We have tables that contain this information for each industry separately.',\n",
      "    sql=\"```sql\\nSELECT DISTINCT betaChina.Beta_ AS Beta_Value, betaChina.Number_of_firms AS Number_of_Firms\\nFROM betaChina\\nWHERE betaChina.Industry_Name IN ('Software (Entertainment)', 'Software (Internet)', 'Semiconductor')\\nUNION\\nSELECT DISTINCT EVAChina.Beta AS Beta_Value, EVAChina.Number_of_Firms AS Number_of_Firms\\nFROM EVAChina\\nWHERE EVAChina.Industry_Name IN ('Software (Entertainment)', 'Software (Internet)', 'Semiconductor')\\n```\"\n",
      ")\n",
      "Extracted rows:  Beta_Value = 0.9517387989546392, Number_of_Firms = 19\n",
      " Beta_Value = 1.3766066410015918, Number_of_Firms = 173\n",
      " Beta_Value = 1.9284684177392413, Number_of_Firms = 33\n",
      "\n",
      "Beta value and number of firms for the Software (Entertainment) industry are Beta_Value = 0.9517387989546392, Number_of_Firms = 19. For the Software (Internet) industry, the values are Beta_Value = 1.3766066410015918, Number_of_Firms = 173. And for the Semiconductor industry, the values are Beta_Value = 1.9284684177392413, Number_of_Firms = 33.\n"
     ]
    }
   ],
   "source": [
    "tsql = TextToSQLQueryModule(\"China\")\n",
    "sq = tsql(question=\"What is the beta value and number of firms for all the Software industries and semiconductor industry?\")\n",
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beta value and number of firms for the Software (Entertainment) industry are Beta_Value = 0.9517387989546392, Number_of_Firms = 19. For the Software (Internet) industry, the values are Beta_Value = 1.3766066410015918, Number_of_Firms = 173. And for the Semiconductor industry, the values are Beta_Value = 1.9284684177392413, Number_of_Firms = 33.\n"
     ]
    }
   ],
   "source": [
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: taxrate \n",
      "/* \n",
      "col : Industry_name (VARCHAR) | Number_of_firms (INTEGER) | Total_Taxable_Income (INTEGER) | Total_Taxes_Paid_Accrual_ (INTEGER) | Total_Cash_Taxes_Paid (INTEGER) | Cash_Taxes_Accrual_Taxes (INTEGER) | Average_across_all_companies (INTEGER) | Average_across_only_money_making_companies (INTEGER) | Aggregate_tax_rate (INTEGER) | Average_across_only_money_making_companies2 (INTEGER) | Aggregate_tax_rate3 (INTEGER)\n",
      "row 1 : Healthcare Products | 230 | 26657.20900000001 | 5486.0970000000025 | 7358.37 | 1.341275956294611 | 0.0481441302819232 | 0.2058016276197557 | 0.2473315283348905 | 0.2760367748926752 | 0.35300275681386\n",
      "row 2 : Hospitals/Healthcare Facilities | 32 | 12803.409999999998 | 2598.52 | 2117.581 | 0.8149181072302696 | 0.0685657680171947 | 0.2029553064378943 | 0.2241515389301438 | 0.1653919541747081 | 0.18582584339767\n",
      "row 3 : Healthcare Support Services | 119 | 76367.83500000002 | 17205.211 | 18119.770000000008 | 1.0531559304910594 | 0.0808427903603178 | 0.2252939473798097 | 0.2520069365480554 | 0.2372696567867873 | 0.2683567809591\n",
      "row 4 : Heathcare Information and Technology | 128 | 21838.42000000001 | 2553.3450000000003 | 6092.359 | 2.386030481583961 | 0.0311109029990876 | 0.1169198595869114 | 0.2128897982354291 | 0.2789743488768875 | 0.5756510369422\n",
      "row 5 : Total Market (without financials) | 5214 | 2091519.969999999 | 405892.603 | 448383.9540000016 | 1.1046861920762858 | 0.0680489965490159 | 0.1940658510661986 | 0.2166068386944213 | 0.2143818660263625 | 0.25088890170136\n",
      "*/\n",
      "\n",
      "Table name: margin \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Gross_Margin (INTEGER) | Net_Margin (INTEGER) | Pre_tax_Pre_stock_compensation_Operating_Margin (INTEGER) | Pre_tax_Unadjusted_Operating_Margin (INTEGER) | After_tax_Unadjusted_Operating_Margin (INTEGER) | Pre_tax_Lease_adjusted_Margin (INTEGER) | After_tax_Lease_Adjusted_Margin (INTEGER) | Pre_tax_Lease_R_D_adj_Margin (INTEGER) | After_tax_Lease_R_D_adj_Margin (INTEGER) | EBITDA_Sales (INTEGER) | EBITDASG_A_Sales (INTEGER) | EBITDAR_D_Sales (INTEGER) | COGS_Sales (INTEGER) | R_D_Sales (INTEGER) | SG_A_Sales (INTEGER) | Stock_Based_Compensation_Sales (INTEGER) | Lease_Expense_Sales (INTEGER)\n",
      "row 1 : Healthcare Products | 230 | 0.5563922869375535 | 0.0818669682093114 | 0.1558110098462976 | 0.1329496754079793 | 0.1265489289141982 | 0.1339775386003306 | 0.1275273065271049 | 0.1383573960583481 | 0.1319071639851224 | 0.209192854485968 | 0.5282884406553701 | 0.294308887098031 | 0.4436077130624465 | 0.0851160326120629 | 0.3190955861694019 | 0.0228613344383182 | 0.00977793897038\n",
      "row 2 : Hospitals/Healthcare Facilities | 32 | 0.35988660694384 | 0.0512378794350515 | 0.120654788689701 | 0.1156935415013266 | 0.107760924973659 | 0.1154810689328797 | 0.1075630207500502 | 0.1154805210386438 | 0.1075624728558143 | 0.1535618716018076 | 0.1846309639990737 | 0.1535712434768955 | 0.64011339305616 | 9.371875087928916e-06 | 0.031069092397266 | 0.0049612471883743 | 0.02033043025564\n",
      "row 3 : Healthcare Support Services | 119 | 0.1436810802508208 | 0.0225465585118185 | 0.0418808661001017 | 0.0391834198903671 | 0.0360157228905698 | 0.0386507190050279 | 0.0355260870312289 | 0.0386351017430003 | 0.0355104697692013 | 0.0461573478838079 | 0.1419758713498951 | 0.0469153240204038 | 0.8563189197491792 | 0.0007579761365958 | 0.0958185234660871 | 0.0026974462097345 | 0.0040901341131\n",
      "row 4 : Heathcare Information and Technology | 128 | 0.4766949090992926 | 0.0572239502258181 | 0.1655618579641873 | 0.1400806639108432 | 0.1357226279638652 | 0.1403673029983578 | 0.1360003494505323 | 0.1414922855264892 | 0.1371253319786638 | 0.2365968665611617 | 0.4859769860183378 | 0.3029234453206219 | 0.5233050909007073 | 0.0663265787594602 | 0.249380119457176 | 0.0254811940533441 | 0.01221262918932\n",
      "row 5 : Total Market (without financials) | 5214 | 0.3341124303745417 | 0.0759000807387551 | 0.1349854587183475 | 0.1203243079012881 | 0.1121363594881506 | 0.1193029155861628 | 0.1111844718951524 | 0.122305912015432 | 0.1141874683244217 | 0.1649863910103977 | 0.308807292459029 | 0.2037377626528898 | 0.6658875696254583 | 0.038751371642492 | 0.1438209014486313 | 0.0146611508170593 | 0.01534583398930\n",
      "*/\n",
      "\n",
      "Table name: wacc \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | Cost_of_Equity (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER) | Cost_of_Capital (INTEGER) | Cost_of_Capital_Local_Currency_ (INTEGER)\n",
      "row 1 : Healthcare Products | 230 | 1.062079895442844 | 0.0876556751903708 | 0.8876027024944009 | 0.5220553450431473 | 0.053524 | 0.0481441302819232 | 0.040143 | 0.1123972975055991 | 0.0823153789017118 | 0.08231537890171\n",
      "row 2 : Hospitals/Healthcare Facilities | 32 | 0.879576324903561 | 0.0792605109455638 | 0.5563496280559845 | 0.4632806494431785 | 0.050855 | 0.0685657680171947 | 0.03814125 | 0.4436503719440154 | 0.0610179355330013 | 0.06101793553300\n",
      "row 3 : Healthcare Support Services | 119 | 1.0326229061406076 | 0.0863006536824679 | 0.7882476948159491 | 0.495257334489641 | 0.050855 | 0.0808427903603178 | 0.03814125 | 0.2117523051840508 | 0.076102788936416 | 0.0761027889364\n",
      "row 4 : Heathcare Information and Technology | 128 | 1.2747542568245518 | 0.0974386958139293 | 0.8615496523681622 | 0.5415316524692178 | 0.053524 | 0.0311109029990876 | 0.040143 | 0.1384503476318378 | 0.0895060868106828 | 0.08950608681068\n",
      "row 5 : Total Market (without financials) | 5214 | 1.097522642247185 | 0.0892860415433705 | 0.815661408370401 | 0.4700345950094993 | 0.050855 | 0.0680489965490159 | 0.03814125 | 0.1843385916295989 | 0.0798580827010761 | 0.07985808270107\n",
      "*/\n",
      "\n",
      "Table name: histgr \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | CAGR_in_Net_Income_Last_5_years (INTEGER) | CAGR_in_Revenues_Last_5_years (INTEGER) | Expected_Growth_in_Revenues_Next_2_years (INTEGER) | Expected_growth_in_Revenues_Next_5_years (INTEGER) | Expected_Growth_in_EPS_Next_5_years (INTEGER)\n",
      "row 1 : Healthcare Products | 230 | 0.1271279069767442 | 0.1794840000000001 | 0.5909131578947368 | 0.2743437920966664 | 0.16686195121951\n",
      "row 2 : Hospitals/Healthcare Facilities | 32 | 0.0409985714285714 | 0.181068947368421 | 0.0413333333333333 | 0.0408655475254498 | 0.10645555555555\n",
      "row 3 : Healthcare Support Services | 119 | 0.0625710344827586 | 0.1242187719298245 | 0.0938797014925373 | 0.1106991073675205 | 0.11186538461538\n",
      "row 4 : Heathcare Information and Technology | 128 | 0.2291124999999999 | 0.1966670689655173 | 0.1482884810126581 | 0.1667594853751504 | 0.10272187499999\n",
      "row 5 : Total Market (without financials) | 5214 | 0.1008838061178445 | 0.1270981652178585 | 0.2757367288045685 | 0.2849241661766458 | 0.14538921052109\n",
      "*/\n",
      "\n",
      "Table name: roe \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | ROE_unadjusted_ (INTEGER) | ROE_adjusted_for_R_D_ (INTEGER)\n",
      "row 1 : Healthcare Products | 230 | 0.0830579792856188 | 0.06980113651478\n",
      "row 2 : Hospitals/Healthcare Facilities | 32 | 0.6196567951123925 | 0.61942896417013\n",
      "row 3 : Healthcare Support Services | 119 | 0.1562031003418411 | 0.1536389472830\n",
      "row 4 : Heathcare Information and Technology | 128 | 0.0517765805634577 | 0.04477886360563\n",
      "row 5 : Total Market (without financials) | 5214 | 0.1659255247192018 | 0.13839857100630\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to calculate the average tax rate of all healthcare industries. To do this, we need to join the taxrate table with the industry_name column as the common key.',\n",
      "    sql=\"```sql\\nSELECT AVG(taxrate.Aggregate_tax_rate) AS Average_Tax_Rate\\nFROM taxrate\\nWHERE taxrate.Industry_name LIKE '%Healthcare%'\\n```\"\n",
      ")\n",
      "Extracted rows:  Average_Tax_Rate = 0.24116333460436337\n",
      "\n",
      "The average tax rate of all healthcare industries is 0.24116333460436337.\n"
     ]
    }
   ],
   "source": [
    "tsql = TextToSQLQueryModule(\"US\")\n",
    "sq = tsql(question=\"What is the average tax rate of all healthcare industries?\")\n",
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The average tax rate of all healthcare industries is 0.24116333460436337.\n"
     ]
    }
   ],
   "source": [
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: taxrate \n",
      "/* \n",
      "col : Industry_name (VARCHAR) | Number_of_firms (INTEGER) | Total_Taxable_Income (INTEGER) | Total_Taxes_Paid_Accrual_ (INTEGER) | Total_Cash_Taxes_Paid (INTEGER) | Cash_Taxes_Accrual_Taxes (INTEGER) | Average_across_all_companies (INTEGER) | Average_across_only_money_making_companies (INTEGER) | Aggregate_tax_rate (INTEGER) | Average_across_only_money_making_companies2 (INTEGER) | Aggregate_tax_rate3 (INTEGER)\n",
      "row 1 : Hospitals/Healthcare Facilities | 32 | 12803.409999999998 | 2598.52 | 2117.581 | 0.8149181072302696 | 0.0685657680171947 | 0.2029553064378943 | 0.2241515389301438 | 0.1653919541747081 | 0.18582584339767\n",
      "row 2 : Healthcare Products | 230 | 26657.20900000001 | 5486.0970000000025 | 7358.37 | 1.341275956294611 | 0.0481441302819232 | 0.2058016276197557 | 0.2473315283348905 | 0.2760367748926752 | 0.35300275681386\n",
      "row 3 : Healthcare Support Services | 119 | 76367.83500000002 | 17205.211 | 18119.770000000008 | 1.0531559304910594 | 0.0808427903603178 | 0.2252939473798097 | 0.2520069365480554 | 0.2372696567867873 | 0.2683567809591\n",
      "row 4 : Heathcare Information and Technology | 128 | 21838.42000000001 | 2553.3450000000003 | 6092.359 | 2.386030481583961 | 0.0311109029990876 | 0.1169198595869114 | 0.2128897982354291 | 0.2789743488768875 | 0.5756510369422\n",
      "row 5 : Total Market (without financials) | 5214 | 2091519.969999999 | 405892.603 | 448383.9540000016 | 1.1046861920762858 | 0.0680489965490159 | 0.1940658510661986 | 0.2166068386944213 | 0.2143818660263625 | 0.25088890170136\n",
      "*/\n",
      "\n",
      "Table name: Employee \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Number_of_Employees (INTEGER) | Market_Capitalization_millions_ (INTEGER) | Revenues_millions_ (INTEGER) | Mkt_Cap_per_Employee_ (INTEGER) | Revenues_per_Employee_ (INTEGER) | Stock_based_Compensation_millions_ (INTEGER) | Stock_based_Compensation_as_of_Revenue (INTEGER)\n",
      "row 1 : Hospitals/Healthcare Facilities | 32 | 716231 | 6704.900000000001 | 7793.799999999999 | 9361.365257856754 | 10881.684819562402 | 688.1889999999999 | 0.00496124718837\n",
      "row 2 : Healthcare Products | 230 | 692739 | 132639.79999999996 | 34790.487 | 191471.5354556333 | 50221.63758645031 | 5056.955000000004 | 0.02286133443831\n",
      "row 3 : Healthcare Support Services | 119 | 1634485 | 55806.99999999999 | 277699.3 | 34143.47638552816 | 169900.18262633184 | 6016.255999999999 | 0.00269744620973\n",
      "row 4 : Heathcare Information and Technology | 128 | 535393 | 104368.758 | 29258.17 | 194938.59277203845 | 54648.02490880531 | 3955.762000000001 | 0.02548119405334\n",
      "row 5 : Total Market (without financials) | 5214 | 37972621 | 10850254.874000004 | 3275178.419000001 | 285738.8978759197 | 86251.04964442673 | 274548.627 | 0.01466115081705\n",
      "*/\n",
      "\n",
      "Table name: margin \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Gross_Margin (INTEGER) | Net_Margin (INTEGER) | Pre_tax_Pre_stock_compensation_Operating_Margin (INTEGER) | Pre_tax_Unadjusted_Operating_Margin (INTEGER) | After_tax_Unadjusted_Operating_Margin (INTEGER) | Pre_tax_Lease_adjusted_Margin (INTEGER) | After_tax_Lease_Adjusted_Margin (INTEGER) | Pre_tax_Lease_R_D_adj_Margin (INTEGER) | After_tax_Lease_R_D_adj_Margin (INTEGER) | EBITDA_Sales (INTEGER) | EBITDASG_A_Sales (INTEGER) | EBITDAR_D_Sales (INTEGER) | COGS_Sales (INTEGER) | R_D_Sales (INTEGER) | SG_A_Sales (INTEGER) | Stock_Based_Compensation_Sales (INTEGER) | Lease_Expense_Sales (INTEGER)\n",
      "row 1 : Hospitals/Healthcare Facilities | 32 | 0.35988660694384 | 0.0512378794350515 | 0.120654788689701 | 0.1156935415013266 | 0.107760924973659 | 0.1154810689328797 | 0.1075630207500502 | 0.1154805210386438 | 0.1075624728558143 | 0.1535618716018076 | 0.1846309639990737 | 0.1535712434768955 | 0.64011339305616 | 9.371875087928916e-06 | 0.031069092397266 | 0.0049612471883743 | 0.02033043025564\n",
      "row 2 : Healthcare Products | 230 | 0.5563922869375535 | 0.0818669682093114 | 0.1558110098462976 | 0.1329496754079793 | 0.1265489289141982 | 0.1339775386003306 | 0.1275273065271049 | 0.1383573960583481 | 0.1319071639851224 | 0.209192854485968 | 0.5282884406553701 | 0.294308887098031 | 0.4436077130624465 | 0.0851160326120629 | 0.3190955861694019 | 0.0228613344383182 | 0.00977793897038\n",
      "row 3 : Healthcare Support Services | 119 | 0.1436810802508208 | 0.0225465585118185 | 0.0418808661001017 | 0.0391834198903671 | 0.0360157228905698 | 0.0386507190050279 | 0.0355260870312289 | 0.0386351017430003 | 0.0355104697692013 | 0.0461573478838079 | 0.1419758713498951 | 0.0469153240204038 | 0.8563189197491792 | 0.0007579761365958 | 0.0958185234660871 | 0.0026974462097345 | 0.0040901341131\n",
      "row 4 : Heathcare Information and Technology | 128 | 0.4766949090992926 | 0.0572239502258181 | 0.1655618579641873 | 0.1400806639108432 | 0.1357226279638652 | 0.1403673029983578 | 0.1360003494505323 | 0.1414922855264892 | 0.1371253319786638 | 0.2365968665611617 | 0.4859769860183378 | 0.3029234453206219 | 0.5233050909007073 | 0.0663265787594602 | 0.249380119457176 | 0.0254811940533441 | 0.01221262918932\n",
      "row 5 : Total Market (without financials) | 5214 | 0.3341124303745417 | 0.0759000807387551 | 0.1349854587183475 | 0.1203243079012881 | 0.1121363594881506 | 0.1193029155861628 | 0.1111844718951524 | 0.122305912015432 | 0.1141874683244217 | 0.1649863910103977 | 0.308807292459029 | 0.2037377626528898 | 0.6658875696254583 | 0.038751371642492 | 0.1438209014486313 | 0.0146611508170593 | 0.01534583398930\n",
      "*/\n",
      "\n",
      "Table name: DollarUS \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Average_Company_Age_years_ (INTEGER) | Market_Cap_millions_ (INTEGER) | Book_Equity_millions_ (INTEGER) | Enteprise_Value_millions_ (INTEGER) | Invested_Capital_millions_ (INTEGER) | Total_Debt_including_leases_millions_ (INTEGER) | Revenues_millions_ (INTEGER) | Gross_Profit_millions_ (INTEGER) | EBITDA_millions_ (INTEGER) | EBIT_Operating_Income_millions_ (INTEGER) | Net_Income_millions_ (INTEGER)\n",
      "row 1 : Hospitals/Healthcare Facilities | 32 | 25.23333333333333 | 122470.973 | 11469.825000000004 | 216360.12943388568 | 87395.24268778635 | 97662.13543388566 | 138712.903 | 49920.916000000005 | 24233.965 | 16018.714313222865 | 7107.3549999999\n",
      "row 2 : Healthcare Products | 230 | 25.382075471698112 | 1050360.2549999994 | 218029.299 | 1141634.0060547118 | 175743.51313973693 | 133007.31705471186 | 221201.21699999992 | 123074.651 | 47150.65200000003 | 29635.99458905761 | 18109.072999999\n",
      "row 3 : Healthcare Support Services | 119 | 28.918367346938776 | 1172283.046 | 321931.98400000005 | 1359335.1278580544 | 189878.4492234088 | 314918.3168580551 | 2230352.538 | 320459.4619999998 | 114747.873 | 86204.72922838898 | 50286.774000000\n",
      "row 4 : Heathcare Information and Technology | 128 | 22.776785714285715 | 738408.3160000001 | 171575.33200000005 | 821060.925284837 | 123028.59595086362 | 118661.63228483696 | 155242.41099999996 | 74003.267 | 34169.073000000004 | 21790.95854303259 | 8883.5840000000\n",
      "row 5 : Total Market (without financials) | 5214 | 34.79463487696991 | 43163609.00299983 | 8566043.05699999 | 51078189.49081551 | 13583012.82046062 | 9754928.713815568 | 18726267.155 | 6256678.631000005 | 3370108.721999998 | 2234098.269636899 | 1421325.1889999\n",
      "*/\n",
      "\n",
      "Table name: fundgrEB \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | ROC (INTEGER) | Reinvestment_Rate (INTEGER) | Expected_Growth_in_EBIT (INTEGER)\n",
      "row 1 : Hospitals/Healthcare Facilities | 32 | 0.2015664022610945 | 0.3491979818825059 | 0.07038658088489\n",
      "row 2 : Healthcare Products | 230 | 0.1308118160288834 | 0.4951245525175228 | 0.06476814187530\n",
      "row 3 : Healthcare Support Services | 119 | 0.4574102830620488 | 0.6275372045055494 | 0.28704197034485\n",
      "row 4 : Heathcare Information and Technology | 128 | 0.1428121724863794 | 0.3133945551322462 | 0.04475655726383\n",
      "row 5 : Total Market (without financials) | 5214 | 0.1466884216172786 | 0.4437349493666002 | 0.06509077933900\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to first filter out the healthcare industries where revenues per employee is more than 1 million, then calculate the average tax rate for those industries.',\n",
      "    sql='```sql\\nSELECT AVG(taxrate.Aggregate_tax_rate) AS Average_Tax_Rate\\nFROM taxrate\\nJOIN Employee ON taxrate.Industry_name = Employee.Industry_Name\\nWHERE Employee.Revenues_per_Employee_ > 1000000\\n```'\n",
      ")\n",
      "Extracted rows:  Average_Tax_Rate = 0.18875589135103757\n",
      "\n",
      "The average tax rate of all healthcare industries where revenues per employee is more than 1 million is approximately 18.88%.\n"
     ]
    }
   ],
   "source": [
    "sq = tsql(question=\"Give me the average tax rate of all healthcare industries where revenues per employee is more than 1 million?\")\n",
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTable name: taxrateEurope \n",
      "/* \n",
      "col : Industry_name (VARCHAR) | Number_of_firms (INTEGER) | Total_Taxable_Income (INTEGER) | Total_Taxes_Paid_Accrual_ (INTEGER) | Total_Cash_Taxes_Paid (INTEGER) | Cash_Taxes_Accrual_Taxes (INTEGER) | Average_across_all_companies (INTEGER) | Average_across_only_money_making_companies (INTEGER) | Aggregate_tax_rate (INTEGER) | Average_across_only_money_making_companies_1 (INTEGER) | Aggregate_tax_rate_1 (INTEGER)\n",
      "row 1 : Financial Svcs. (Non-bank & Insurance) | 143 | 67531.72600000002 | 5859.747000000001 | 4199.801999999999 | 0.716720704835891 | 0.133380349433826 | 0.0867702833480074 | 0.0851491015690155 | 0.0621900586399938 | 0.0637353624270\n",
      "row 2 : Banks (Regional) | 72 | 7968.15 | 1353.5790000000002 | 675.7199999999998 | 0.499209872493589 | 0.1876823623247109 | 0.169873684606841 | 0.1693503510852582 | 0.0848026204325972 | 0.08480262043259\n",
      "row 3 : Total Market (without financials) | 6149 | 1358894.4370000027 | 346515.07499999995 | 347363.15499999904 | 1.0024474548473803 | 0.114021587002736 | 0.2549977875875278 | 0.2805680522510425 | 0.2556218831588302 | 0.29150994341615\n",
      "row 4 : Brokerage & Investment Banking | 74 | 3648.079999999999 | 709.6220000000001 | 633.1540000000001 | 0.8922412213826517 | 0.1518944442193559 | 0.194519308787088 | 0.195143715024721 | 0.1735581456547006 | 0.1756212164289\n",
      "row 5 : Bank (Money Center) | 112 | 300482.43000000005 | 70812.87999999998 | 18826.331 | 0.2658602644038769 | 0.2354411572737627 | 0.2356639621158547 | 0.2359508773550057 | 0.0626536832785863 | 0.06272274286469\n",
      "*/\n",
      "\n",
      "Table name: betaEurope \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Beta_ (INTEGER) | D_E_Ratio (INTEGER) | Effective_Tax_rate (INTEGER) | Unlevered_beta (INTEGER) | Cash_Firm_value (INTEGER) | Unlevered_beta_corrected_for_cash (INTEGER) | HiLo_Risk (INTEGER) | Standard_deviation_of_equity (INTEGER) | Standard_deviation_in_operating_income_last_10_years_ (INTEGER) | 2020 (INTEGER) | 2021_0 (INTEGER) | 2022_0 (INTEGER) | 2022_0_1 (INTEGER) | Average_2020_24 (INTEGER)\n",
      "row 1 : Financial Svcs. (Non-bank & Insurance) | 143 | 1.0216367356825768 | 4.109932277910462 | 0.133380349433826 | 0.2495224495401071 | 0.2851940249654809 | 0.3490771737436264 | 0.3644714846498691 | 0.355513616932455 | 0.4724600661983447 | 0.26 | 0.27 | 0.33 | 0.2491729799474152 | 0.29165003073820\n",
      "row 2 : Total Market (without financials) | 6149 | 1.0163277560226116 | 0.3625130044987187 | 0.114021587002736 | 0.7984122714874734 | 0.0654744498128502 | 0.8543503934457243 | 0.3678598807686726 | 0.3394865864311263 | 0.2520005799939234 | 0.85 | 0.78 | 0.87 | 0.8817115657835057 | 0.8472123918458\n",
      "row 3 : Brokerage & Investment Banking | 74 | 0.8707315899058546 | 0.6876389626893673 | 0.1518944442193559 | 0.573709019867202 | 0.2412741526659604 | 0.7561479839958829 | 0.3177625193382241 | 0.2926244614542751 | 0.8202608076286918 | 0.95 | 0.73 | 0.9 | 0.6959319603512203 | 0.80641598886942\n",
      "row 4 : Bank (Money Center) | 112 | 1.2153935486343157 | 5.615675883136155 | 0.2354411572737627 | 0.2324758412532342 | 0.3625662794009818 | 0.3647059039091448 | 0.2449799817999482 | 0.2535837672757894 | 1.5390713693641629 | 0.32 | 0.28 | 0.42 | 0.3686115410275469 | 0.35066348898733\n",
      "row 5 : Retail (REITs) | 33 | 0.960746998690458 | 1.4416481088663338 | 0.0421971427092918 | 0.4606978185423953 | 0.0706780075754195 | 0.4957354095757971 | 0.1860296862232068 | 0.2120972268727357 | 0.2453273616689489 | 1.23 | 1.12 | 1.25 | 1.4672850705706748 | 1.11260409602929\n",
      "*/\n",
      "\n",
      "Table name: dbtfundEurope \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_firms (INTEGER) | Book_Debt_to_Capital (INTEGER) | Market_Debt_to_Capital_Unadjusted_ (INTEGER) | Market_D_E_unadjusted_ (INTEGER) | Market_Debt_to_Capital_adjusted_for_leases_ (INTEGER) | Market_D_E_adjusted_for_leases_ (INTEGER) | Interest_Coverage_Ratio (INTEGER) | Debt_to_EBITDA (INTEGER) | Effective_tax_rate (INTEGER) | Institutional_Holdings (INTEGER) | Std_dev_in_Stock_Prices (INTEGER) | EBITDA_EV (INTEGER) | Net_PP_E_Total_Assets (INTEGER) | Capital_Spending_Total_Assets (INTEGER)\n",
      "row 1 : Financial Svcs. (Non-bank & Insurance) | 143 | 0.8064776202399659 | 0.8042998495207718 | 4.109858104611632 | 0.8043026902092493 | 4.109932277910462 | 5.282640237538722 | 134.02896372206712 | 0.133380349433826 | 0.2892209174311926 | 0.355513616932455 | 0.0106491195890414 | 0.0051764605692062 | 0.00082237262437\n",
      "row 2 : Total Market (without financials) | 6149 | 0.4398362396046586 | 0.2654959269193548 | 0.3614628381920554 | 0.2660620510055906 | 0.3625130044987187 | 7.095364720701234 | 2.8548850784071864 | 0.114021587002736 | 0.2515864670455792 | 0.3395466680004999 | 0.1015165247503355 | 0.2846296929541041 | 0.03637768238513\n",
      "row 3 : Brokerage & Investment Banking | 74 | 0.6232840556024103 | 0.4073341226467672 | 0.6872913359984001 | 0.4074562023583337 | 0.6876389626893673 | 2.863783019167806 | 60.495499496601965 | 0.1518944442193559 | 0.2122024561403509 | 0.2926244614542751 | 0.0079760754842969 | 0.0115649007081163 | 0.00158586948421\n",
      "row 4 : Bank (Money Center) | 112 | 0.8088976777082401 | 0.8488435545278178 | 5.615662315134511 | 0.8488438645325607 | 5.615675883136155 | -0.1518602885345482 | -965121.437387172 | 0.2354411572737627 | 0.2839601801801802 | 0.2535837672757894 | 7.41793318300574e-07 | 0.0097046332187522 | 0.00033102263072\n",
      "row 5 : Retail (REITs) | 33 | 0.4885361986828798 | 0.5904405308196344 | 1.441647856418162 | 0.5904405731650234 | 1.4416481088663338 | 2.737197146373365 | 12.669698049396471 | 0.0421971427092918 | 0.3597851724137931 | 0.2120972268727357 | 0.0484124926581707 | 0.7223189363320469 | 0.00068097022051\n",
      "*/\n",
      "\n",
      "Table name: fundgrEBEurope \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | ROC (INTEGER) | Reinvestment_Rate (INTEGER) | Expected_Growth_in_EBIT (INTEGER)\n",
      "row 1 : Financial Svcs. (Non-bank & Insurance) | 158 | 0.0090540861074978 | -12.65841678108797 | -0.11461039552056\n",
      "row 2 : Banks (Regional) | 72 | 2.806938756766444e-06 | 1149.8288982874033 | 0.00322749929825\n",
      "row 3 : Total Market (without financials) | 6134 | 0.1446756721246595 | 0.3207162519218786 | 0.04639983930809\n",
      "row 4 : Brokerage & Investment Banking | 74 | 0.0115528682067912 | 1.892566684166849 | 0.02186457347474\n",
      "row 5 : Retail (REITs) | 33 | 0.0428673616720445 | 0.0452603514308576 | 0.00194019185419\n",
      "*/\n",
      "\n",
      "Table name: waccEurope \n",
      "/* \n",
      "col : Industry_Name (VARCHAR) | Number_of_Firms (INTEGER) | Beta (INTEGER) | Cost_of_Equity (INTEGER) | E_D_E_ (INTEGER) | Std_Dev_in_Stock (INTEGER) | Cost_of_Debt (INTEGER) | Tax_Rate (INTEGER) | After_tax_Cost_of_Debt (INTEGER) | D_D_E_ (INTEGER) | Cost_of_Capital (INTEGER) | Cost_of_Capital_Euros_ (INTEGER)\n",
      "row 1 : Financial Svcs. (Non-bank & Insurance) | 143 | 1.0216367356825768 | 0.0989744037317037 | 0.1956973097907507 | 0.355513616932455 | 0.0604549999999999 | 0.133380349433826 | 0.0455165695 | 0.8043026902092493 | 0.0559781238463843 | 0.04059980165444\n",
      "row 2 : Banks (Regional) | 72 | 0.4682886846720631 | 0.0663822035271845 | 0.2580269829998616 | 0.132835433574351 | 0.054643 | 0.1876823623247109 | 0.0411407147 | 0.7419730170001384 | 0.0476536999085031 | 0.03239660719138\n",
      "row 3 : Total Market (without financials) | 6149 | 1.0163277560226116 | 0.0986617048297318 | 0.7339379489944093 | 0.3395466680004999 | 0.0604549999999999 | 0.114021587002736 | 0.0455165695 | 0.2660620510055906 | 0.0845218011229337 | 0.06872779431046\n",
      "row 4 : Brokerage & Investment Banking | 74 | 0.8707315899058546 | 0.0900860906454548 | 0.5925437976416663 | 0.2926244614542751 | 0.0604549999999999 | 0.1518944442193559 | 0.0455165695 | 0.4074562023583337 | 0.0719259628185983 | 0.0563153905445\n",
      "row 5 : Bank (Money Center) | 112 | 1.2153935486343157 | 0.1103866800145611 | 0.1511561354674393 | 0.2535837672757894 | 0.0604549999999999 | 0.2354411572737627 | 0.0455165695 | 0.8488438645325607 | 0.0553220847127267 | 0.03995331648875\n",
      "*/\n",
      "\n",
      "\u001b[0m\n",
      "Prediction(\n",
      "    rationale='produce the SQL. We need to find the average tax rate of all banking industries where the number of firms is more than 500. To do this, we need to join the taxrateEurope table with the betaEurope table on the Industry_Name column and filter the results based on the Number_of_firms column.',\n",
      "    sql=\"```sql\\nSELECT AVG(te.Aggregate_tax_rate) AS Average_Tax_Rate\\nFROM taxrateEurope te\\nJOIN betaEurope be ON te.Industry_name = be.Industry_Name\\nWHERE te.Industry_name LIKE '%Bank%'\\nAND te.Number_of_firms > 500\\n```\"\n",
      ")\n",
      "Extracted rows:  Average_Tax_Rate = None\n",
      "\n",
      "There are no relevant rows that match the criteria of banking industries with more than 500 firms in the database.\n"
     ]
    }
   ],
   "source": [
    "tsql = TextToSQLQueryModule(\"Europe\")\n",
    "sq = tsql(question=\"Give me the average tax rate of all banking industries where number of firms is more than 500?\")\n",
    "print(sq.answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
